diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 9ca0cf0aaae918..69bc1c4662db0d 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -1565,7 +1565,6 @@ cc_library( "@xla//xla/service:hlo_proto_cc_impl", "@xla//xla/service:metrics_proto_cc_impl", "@xla//xla/service:buffer_assignment_proto_cc_impl", - "@xla//xla/service/memory_space_assignment:memory_space_assignment_proto_cc_impl", "@xla//xla/service/gpu:backend_configs_cc_impl", "@xla//xla/service/gpu/model:hlo_op_profile_proto_cc_impl", "@xla//xla/stream_executor:device_description_proto_cc_impl", diff --git a/tensorflow/core/kernels/ragged_cross_op.cc b/tensorflow/core/kernels/ragged_cross_op.cc index e1731776070e44..bc59f38adbac8d 100644 --- a/tensorflow/core/kernels/ragged_cross_op.cc +++ b/tensorflow/core/kernels/ragged_cross_op.cc @@ -410,6 +410,13 @@ class RaggedCrossOp : public OpKernel { return absl::InvalidArgumentError( "tf.ragged.cross only supports inputs with rank=2."); } + if (sparse_indices_list[i].shape().dim_size(0) != + sparse_values_list[i].shape().dim_size(0)) { + return absl::InvalidArgumentError(absl::StrCat( + "Expected size of values to be ", + sparse_indices_list[i].shape().dim_size(0), " got ", + sparse_values_list[i].shape().dim_size(0), " at position ", i)); + } } for (int i = 0; i < dense_list.size(); ++i) { if (!TensorShapeUtils::IsMatrix(dense_list[i].shape())) { diff --git a/tensorflow/dtensor/python/d_random.py b/tensorflow/dtensor/python/d_random.py index 69aacf486a0932..5e2c5686f319f5 100644 --- a/tensorflow/dtensor/python/d_random.py +++ b/tensorflow/dtensor/python/d_random.py @@ -105,7 +105,7 @@ def _old_tf_random_stateless_uniform( with ops.name_scope( name, "stateless_random_uniform", [shape, seed, minval, maxval] ) as name: - seed = ops.convert_to_tensor(seed, dtype_hint=dtypes.int32, name="seed") + seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed") if dtype.is_integer and minval is None and maxval is None: result = api.call_with_layout( @@ -267,7 +267,7 @@ def stateless_split(seed, num=2, mesh=None): def _get_noise_shape(x, noise_shape): - """Noisve shape util copied from tf nn_ops.""" + """Noise shape util copied from tf nn_ops.""" # If noise_shape is none return immediately. if noise_shape is None: return array_ops.shape(x) diff --git a/tensorflow/lite/delegates/flex/kernel.cc b/tensorflow/lite/delegates/flex/kernel.cc index 3faa1b71ea1fd2..28528fd3ed9502 100644 --- a/tensorflow/lite/delegates/flex/kernel.cc +++ b/tensorflow/lite/delegates/flex/kernel.cc @@ -235,7 +235,7 @@ class OpNode { absl::Status InitializeNodeDef(const void* custom_initial_data, int custom_initial_data_size) { if (!custom_initial_data) { - return tensorflow::errors::Internal( + return absl::InternalError( "Cannot convert empty data into a valid NodeDef"); } // The flexbuffer contains a vector where the first elements is the @@ -249,8 +249,7 @@ class OpNode { name_ = v[0].AsString().str(); if (!nodedef_.ParseFromString(v[1].AsString().str())) { nodedef_.Clear(); - return tensorflow::errors::Internal( - "Failed to parse data into a valid NodeDef"); + return absl::InternalError("Failed to parse data into a valid NodeDef"); } // Fill NodeDef with defaults if it's a valid op. @@ -306,8 +305,8 @@ class OpNode { // Lite native buffer, or could be produced by a separater subgraph). We // need to fetch it from the delegate's buffer_map. if (!buffer_map->HasTensor(input_index)) { - return tensorflow::errors::Internal( - "Cannot read from invalid tensor index ", input_index); + return absl::InternalError(absl::StrCat( + "Cannot read from invalid tensor index ", input_index)); } run_state->input_tf_tensors[i] = buffer_map->GetTensor(input_index); } else { diff --git a/tensorflow/lite/delegates/ynnpack/BUILD b/tensorflow/lite/delegates/ynnpack/BUILD index 92a14c4a6ab04b..16315e6414a79a 100644 --- a/tensorflow/lite/delegates/ynnpack/BUILD +++ b/tensorflow/lite/delegates/ynnpack/BUILD @@ -107,7 +107,9 @@ cc_test( "//tensorflow/lite/c:common", "//tensorflow/lite/kernels:test_util", "//tensorflow/lite/schema:schema_fbs", + "//tensorflow/lite/types:half", "@com_google_googletest//:gtest_main", + "@eigen_archive//:eigen3", "@flatbuffers//:runtime_cc", ], ) @@ -124,6 +126,7 @@ cc_binary( "//tensorflow/lite/c:common", "//tensorflow/lite/kernels:test_util", "//tensorflow/lite/schema:schema_fbs", + "@com_google_absl//absl/algorithm:container", "@com_google_benchmark//:benchmark_main", "@flatbuffers//:runtime_cc", ], diff --git a/tensorflow/lite/delegates/ynnpack/moe.cc b/tensorflow/lite/delegates/ynnpack/moe.cc index bc7eb66908dad0..ffc64f4767d2a4 100644 --- a/tensorflow/lite/delegates/ynnpack/moe.cc +++ b/tensorflow/lite/delegates/ynnpack/moe.cc @@ -30,137 +30,275 @@ limitations under the License. namespace tflite { namespace ynnpack { +namespace { + +// The MoE op comes in two input layouts: +// - the base layout: tokens, routing weights, expert indices, gate/up/down +// weights and the output scale. +// - the extended layout, which additionally carries an explicit per-expert +// scale tensor for each of the gate/up/down weights. +constexpr int kNumInputsWithoutWeightScales = 7; +constexpr int kNumInputsWithWeightScales = 10; + +// Positions of the MoE inputs in `node->inputs`. The `*_scale` indices are +// `-1` unless the op uses the extended layout. +struct MoeInputIndices { + int tokens; + int routing_weights; + int expert_indices; + int gate; + int gate_scale; + int up; + int up_scale; + int down; + int down_scale; + int scale; +}; + +MoeInputIndices GetInputIndices(bool has_weight_scales) { + return MoeInputIndices{ + /*tokens=*/0, + /*routing_weights=*/1, + /*expert_indices=*/2, + /*gate=*/3, + /*gate_scale=*/has_weight_scales ? 4 : -1, + /*up=*/has_weight_scales ? 5 : 4, + /*up_scale=*/has_weight_scales ? 6 : -1, + /*down=*/has_weight_scales ? 7 : 5, + /*down_scale=*/has_weight_scales ? 8 : -1, + /*scale=*/has_weight_scales ? 9 : 6, + }; +} + +bool Is16BitFloatType(TfLiteType type) { + return type == kTfLiteFloat16 || type == kTfLiteBFloat16; +} + +bool IsFloatType(TfLiteType type) { + return type == kTfLiteFloat32 || Is16BitFloatType(type); +} + +// Weight types that make the MoE subgraph run in dynamically quantized mode. +bool IsQuantizedWeightType(TfLiteType type) { + return type == kTfLiteInt8 || type == kTfLiteInt4 || type == kTfLiteInt2; +} + +// Converts `value_id` to `to`, or leaves it alone if it already has that type. +TfLiteStatus ConvertIfNeeded(TfLiteContext* context, ynn_subgraph_t subgraph, + TfLiteType from, TfLiteType to, + uint32_t& value_id) { + if (from == to) { + return kTfLiteOk; + } + uint32_t converted = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS( + ynn_define_convert(subgraph, value_id, GetYnnType(to), &converted, 0)); + value_id = converted; + return kTfLiteOk; +} + +// One of the gate/up/down weight tensors, together with the graph values +// derived from it. +struct WeightValues { + int tensor_index; + // Index of the explicit per-expert scale tensor, or `-1` if the op does not + // carry one. + int scale_tensor_index; + // Weights, after transposition and gathering of the active experts. + uint32_t value_id = YNN_INVALID_VALUE_ID; + // Quantization parameters, invalid for float weights. + uint32_t scale_id = YNN_INVALID_VALUE_ID; + uint32_t zero_point_id = YNN_INVALID_VALUE_ID; +}; + +// The three weight tensors, in the order they appear in `node->inputs` and in +// the order the subgraph consumes them. +enum WeightKind { kGate = 0, kUp = 1, kDown = 2 }; +constexpr int kNumWeights = 3; + +} // namespace bool IsMoe(const TfLiteRegistration* registration, const TfLiteNode* node) { - if (registration == nullptr) { + if (registration == nullptr || node == nullptr || + registration->builtin_code != kTfLiteBuiltinStablehloComposite || + node->builtin_data == nullptr) { return false; } - if (registration->builtin_code == kTfLiteBuiltinStablehloComposite && - node != nullptr && node->builtin_data != nullptr) { - const auto* composite_params = - static_cast(node->builtin_data); - return composite_params->name != nullptr && - strcmp(composite_params->name, "odml.moe_experts") == 0; - } - return false; + const auto* composite_params = + static_cast(node->builtin_data); + return composite_params->name != nullptr && + strcmp(composite_params->name, "odml.moe_experts") == 0; } TfLiteStatus IsMoeSupported(const TfLiteRegistration* registration, const TfLiteNode* node, TfLiteContext* context) { TF_LITE_ENSURE(context, IsMoe(registration, node)); - TF_LITE_ENSURE_EQ(context, node->inputs->size, 7); + TF_LITE_ENSURE(context, node->inputs->size == kNumInputsWithoutWeightScales || + node->inputs->size == kNumInputsWithWeightScales); TF_LITE_ENSURE_EQ(context, node->outputs->size, 1); - const TfLiteTensor& tokens = context->tensors[node->inputs->data[0]]; - const TfLiteTensor& routing_weights = context->tensors[node->inputs->data[1]]; - const TfLiteTensor& expert_indices = context->tensors[node->inputs->data[2]]; - const TfLiteTensor& gate_weights = context->tensors[node->inputs->data[3]]; - const TfLiteTensor& up_weights = context->tensors[node->inputs->data[4]]; - const TfLiteTensor& down_weights = context->tensors[node->inputs->data[5]]; - const TfLiteTensor& scale = context->tensors[node->inputs->data[6]]; + const bool has_weight_scales = + node->inputs->size == kNumInputsWithWeightScales; + const MoeInputIndices idx = GetInputIndices(has_weight_scales); + + const TfLiteTensor& tokens = context->tensors[node->inputs->data[idx.tokens]]; + const TfLiteTensor& routing_weights = + context->tensors[node->inputs->data[idx.routing_weights]]; + const TfLiteTensor& expert_indices = + context->tensors[node->inputs->data[idx.expert_indices]]; + const TfLiteTensor& gate_weights = + context->tensors[node->inputs->data[idx.gate]]; + const TfLiteTensor& up_weights = context->tensors[node->inputs->data[idx.up]]; + const TfLiteTensor& down_weights = + context->tensors[node->inputs->data[idx.down]]; + const TfLiteTensor& scale = context->tensors[node->inputs->data[idx.scale]]; const TfLiteTensor& output = context->tensors[node->outputs->data[0]]; TF_LITE_ENSURE(context, IsTensorSupported(tokens)); TF_LITE_ENSURE(context, IsTensorSupported(routing_weights)); TF_LITE_ENSURE(context, IsTensorSupported(expert_indices)); - TF_LITE_ENSURE(context, IsTensorSupported(gate_weights)); - TF_LITE_ENSURE(context, IsTensorSupported(up_weights)); - TF_LITE_ENSURE(context, IsTensorSupported(down_weights)); TF_LITE_ENSURE(context, IsTensorSupported(scale)); TF_LITE_ENSURE(context, IsTensorSupported(output)); - TF_LITE_ENSURE(context, tflite::IsConstantTensor(&gate_weights)); - TF_LITE_ENSURE(context, tflite::IsConstantTensor(&up_weights)); - TF_LITE_ENSURE(context, tflite::IsConstantTensor(&down_weights)); TF_LITE_ENSURE(context, tflite::IsConstantTensor(&scale)); - TF_LITE_ENSURE_EQ(context, tokens.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, routing_weights.type, kTfLiteFloat32); + TF_LITE_ENSURE(context, IsFloatType(tokens.type)); + TF_LITE_ENSURE(context, IsFloatType(routing_weights.type)); TF_LITE_ENSURE_EQ(context, expert_indices.type, kTfLiteInt32); - TF_LITE_ENSURE_EQ(context, gate_weights.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, up_weights.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, down_weights.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, scale.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, output.type, kTfLiteFloat32); + TF_LITE_ENSURE(context, IsFloatType(scale.type)); + TF_LITE_ENSURE(context, IsFloatType(output.type)); + TF_LITE_ENSURE_EQ(context, tokens.type, output.type); + + const TfLiteTensor* weights[] = {&gate_weights, &up_weights, &down_weights}; + for (const TfLiteTensor* w : weights) { + TF_LITE_ENSURE(context, IsTensorSupported(*w, /*allow_per_channel=*/true)); + TF_LITE_ENSURE(context, tflite::IsConstantTensor(w)); + TF_LITE_ENSURE(context, + IsFloatType(w->type) || IsQuantizedWeightType(w->type)); + TF_LITE_ENSURE_EQ(context, w->type, gate_weights.type); + TF_LITE_ENSURE(context, w->dims != nullptr && w->dims->size == 4); + } + + if (has_weight_scales) { + const int scale_indices[] = {idx.gate_scale, idx.up_scale, idx.down_scale}; + for (int scale_index : scale_indices) { + const TfLiteTensor& weight_scale = + context->tensors[node->inputs->data[scale_index]]; + TF_LITE_ENSURE(context, IsTensorSupported(weight_scale)); + TF_LITE_ENSURE(context, tflite::IsConstantTensor(&weight_scale)); + TF_LITE_ENSURE(context, IsFloatType(weight_scale.type)); + } + } TF_LITE_ENSURE(context, tokens.dims != nullptr && tokens.dims->size == 3); TF_LITE_ENSURE(context, routing_weights.dims != nullptr && routing_weights.dims->size == 3); TF_LITE_ENSURE(context, expert_indices.dims != nullptr && expert_indices.dims->size == 3); - TF_LITE_ENSURE(context, - gate_weights.dims != nullptr && gate_weights.dims->size == 4); - TF_LITE_ENSURE(context, - up_weights.dims != nullptr && up_weights.dims->size == 4); - TF_LITE_ENSURE(context, - down_weights.dims != nullptr && down_weights.dims->size == 4); TF_LITE_ENSURE(context, output.dims != nullptr && output.dims->size == 3); + // Sub-byte weights pack several values into one byte along the innermost + // dimension. `IsTensorSupported` checks that for the layout the weights + // arrive in, but `DefineMoeNode` transposes [X, E, 1, Y] to [E, Y, X], which + // makes the leading dimension the innermost one, so it has to hold a whole + // number of bytes too. + const int weights_per_byte = + static_cast(YnnTypeElementCount(GetYnnType(gate_weights.type))); + if (weights_per_byte > 1) { + for (const TfLiteTensor* w : weights) { + TF_LITE_ENSURE_EQ(context, w->dims->data[0] % weights_per_byte, 0); + } + } + return kTfLiteOk; } TfLiteStatus DefineMoeNode(TfLiteContext* context, ynn_subgraph_t subgraph, TensorToValueIdMap& tensor_to_value_id, const NodeInfo& node) { - TF_LITE_ENSURE_EQ(context, node.inputs.size(), 7); + TF_LITE_ENSURE(context, node.inputs.size() == kNumInputsWithoutWeightScales || + node.inputs.size() == kNumInputsWithWeightScales); TF_LITE_ENSURE_EQ(context, node.outputs.size(), 1); - const TfLiteTensor& tokens = context->tensors[node.inputs[0]]; - const TfLiteTensor& routing_weights = context->tensors[node.inputs[1]]; - const TfLiteTensor& expert_indices = context->tensors[node.inputs[2]]; - const TfLiteTensor& gate_weights = context->tensors[node.inputs[3]]; - const TfLiteTensor& up_weights = context->tensors[node.inputs[4]]; - const TfLiteTensor& down_weights = context->tensors[node.inputs[5]]; - const TfLiteTensor& scale = context->tensors[node.inputs[6]]; + const bool has_weight_scales = + node.inputs.size() == kNumInputsWithWeightScales; + const MoeInputIndices idx = GetInputIndices(has_weight_scales); + + const TfLiteTensor& tokens = context->tensors[node.inputs[idx.tokens]]; + const TfLiteTensor& routing_weights = + context->tensors[node.inputs[idx.routing_weights]]; + const TfLiteTensor& expert_indices = + context->tensors[node.inputs[idx.expert_indices]]; + const TfLiteTensor& gate_weights = context->tensors[node.inputs[idx.gate]]; + const TfLiteTensor& scale = context->tensors[node.inputs[idx.scale]]; - TF_LITE_ENSURE_EQ(context, tokens.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, routing_weights.type, kTfLiteFloat32); + TF_LITE_ENSURE(context, IsFloatType(tokens.type)); + TF_LITE_ENSURE(context, IsFloatType(routing_weights.type)); TF_LITE_ENSURE_EQ(context, expert_indices.type, kTfLiteInt32); - TF_LITE_ENSURE_EQ(context, gate_weights.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, up_weights.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, down_weights.type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, scale.type, kTfLiteFloat32); + TF_LITE_ENSURE(context, IsFloatType(scale.type)); - int E = gate_weights.dims->data[1]; - int K = (expert_indices.dims && expert_indices.dims->size >= 1) - ? expert_indices.dims->data[expert_indices.dims->size - 1] - : 0; - TF_LITE_ENSURE(context, K > 0); + const int num_experts = gate_weights.dims->data[1]; + const int num_selected_experts = + (expert_indices.dims && expert_indices.dims->size >= 1) + ? expert_indices.dims->data[expert_indices.dims->size - 1] + : 0; + TF_LITE_ENSURE(context, num_selected_experts > 0); TF_LITE_ENSURE(context, tflite::NumElements(&scale) == 1 || - tflite::NumElements(&scale) == E); + tflite::NumElements(&scale) == num_experts); - uint32_t gate_weights_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[3]); - uint32_t up_weights_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[4]); - uint32_t down_weights_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[5]); - uint32_t scale_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[6]); + const bool is_quantized = IsQuantizedWeightType(gate_weights.type); + + // A float dot needs both operands in the same type: YNNPACK has fp32, fp16 + // and bf16 kernels, all of which accumulate in fp32, but none for a mixed + // pair. + // TODO: If YNNPACK adds mixed-type float dot kernels, remove this conversion. + const TfLiteType dot_float_type = + tokens.type == gate_weights.type ? tokens.type : kTfLiteFloat32; + + WeightValues weights[kNumWeights] = { + {node.inputs[idx.gate], + has_weight_scales ? node.inputs[idx.gate_scale] : -1}, + {node.inputs[idx.up], has_weight_scales ? node.inputs[idx.up_scale] : -1}, + {node.inputs[idx.down], + has_weight_scales ? node.inputs[idx.down_scale] : -1}, + }; + + uint32_t scale_val = GetOrCreateValueId(context, subgraph, tensor_to_value_id, + node.inputs[idx.scale]); // 1. Transpose constant weights from 4D [*, E, 1, *] to 3D [E, *, *] // dropping the unit axis 2 so expert dimension E is on axis 0 and can be // gathered directly. int32_t perm_weights[] = {1, 3, 0}; - uint32_t w_gate_transposed = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose( - subgraph, 3, perm_weights, gate_weights_val, &w_gate_transposed, 0)); - - uint32_t w_up_transposed = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose( - subgraph, 3, perm_weights, up_weights_val, &w_up_transposed, 0)); - - uint32_t w_down_transposed = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose( - subgraph, 3, perm_weights, down_weights_val, &w_down_transposed, 0)); - - uint32_t tokens_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[0]); - uint32_t routing_weights_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[1]); - uint32_t expert_indices_val = - GetOrCreateValueId(context, subgraph, tensor_to_value_id, node.inputs[2]); + for (WeightValues& w : weights) { + const uint32_t weights_val = GetOrCreateValueId( + context, subgraph, tensor_to_value_id, w.tensor_index); + TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose( + subgraph, 3, perm_weights, weights_val, &w.value_id, 0)); + } + + if (is_quantized) { + for (WeightValues& w : weights) { + if (has_weight_scales) { + const uint32_t weight_scale_val = GetOrCreateValueId( + context, subgraph, tensor_to_value_id, w.scale_tensor_index); + TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose( + subgraph, 3, perm_weights, weight_scale_val, &w.scale_id, 0)); + } else { + TF_LITE_ENSURE_STATUS(DefineQuantizationParams( + context, subgraph, context->tensors[w.tensor_index], &w.scale_id, + &w.zero_point_id)); + } + } + } + + uint32_t tokens_val = GetOrCreateValueId( + context, subgraph, tensor_to_value_id, node.inputs[idx.tokens]); + uint32_t routing_weights_val = GetOrCreateValueId( + context, subgraph, tensor_to_value_id, node.inputs[idx.routing_weights]); + uint32_t expert_indices_val = GetOrCreateValueId( + context, subgraph, tensor_to_value_id, node.inputs[idx.expert_indices]); uint32_t output_val = YNN_INVALID_VALUE_ID; auto out_it = tensor_to_value_id.find(node.outputs[0]); if (out_it != tensor_to_value_id.end()) { @@ -173,6 +311,11 @@ TfLiteStatus DefineMoeNode(TfLiteContext* context, ynn_subgraph_t subgraph, TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_expand_dims( subgraph, 2, expand_axes_2_3, tokens_val, &tokens_5d, 0)); + if (!is_quantized) { + TF_LITE_ENSURE_STATUS(ConvertIfNeeded(context, subgraph, tokens.type, + dot_float_type, tokens_5d)); + } + // 3. Expand expert indices: [B, N, K] -> [B, N, K, 1, 1] uint32_t ei_5d = YNN_INVALID_VALUE_ID; int32_t expand_axes_3_4[] = {3, 4}; @@ -182,41 +325,125 @@ TfLiteStatus DefineMoeNode(TfLiteContext* context, ynn_subgraph_t subgraph, // 4. Gather weights along axis 0 (expert dim E) for active choices: // [E, D_in, D_mid] gathered by [B, N, K, 1, 1] -> [B, N, K, D_in, D_mid] int32_t gather_axis_0 = 0; - uint32_t w_gate_k = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_gather( - subgraph, 1, &gather_axis_0, 5, w_gate_transposed, ei_5d, &w_gate_k, 0)); + for (WeightValues& w : weights) { + uint32_t gathered = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS(ynn_define_gather( + subgraph, 1, &gather_axis_0, 5, w.value_id, ei_5d, &gathered, 0)); + w.value_id = gathered; + } - uint32_t w_up_k = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_gather( - subgraph, 1, &gather_axis_0, 5, w_up_transposed, ei_5d, &w_up_k, 0)); + if (!is_quantized) { + for (WeightValues& w : weights) { + TF_LITE_ENSURE_STATUS(ConvertIfNeeded( + context, subgraph, context->tensors[w.tensor_index].type, + dot_float_type, w.value_id)); + } + } - uint32_t w_down_k = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_gather( - subgraph, 1, &gather_axis_0, 5, w_down_transposed, ei_5d, &w_down_k, 0)); + if (is_quantized && has_weight_scales) { + for (WeightValues& w : weights) { + uint32_t gathered = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS(ynn_define_gather( + subgraph, 1, &gather_axis_0, 5, w.scale_id, ei_5d, &gathered, 0)); + w.scale_id = gathered; + } + } + + auto dynamic_quantize = [&](uint32_t input_id, uint32_t* quant_id, + uint32_t* zp_id, + uint32_t* scale_id) -> TfLiteStatus { + int32_t reduce_axis_last = -1; + uint32_t min_max_id = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS(ynn_define_reduce( + subgraph, ynn_reduce_min_max, 1, &reduce_axis_last, input_id, + YNN_INVALID_VALUE_ID, &min_max_id, YNN_NODE_FLAG_KEEP_DIMS)); + + TF_LITE_ENSURE_YNN_STATUS(ynn_define_dynamic_quantization( + subgraph, min_max_id, ynn_type_int8, zp_id, scale_id, 0)); + + TF_LITE_ENSURE_YNN_STATUS(ynn_define_quantize( + subgraph, input_id, ynn_type_int8, *zp_id, *scale_id, quant_id, 0)); + return kTfLiteOk; + }; + + // A dot against dynamically quantized activations has to carry both operands' + // quantization parameters and dequantize the int32 accumulator; a float dot + // is just a dot. + auto define_dot = [&](uint32_t a_id, uint32_t a_zp, uint32_t a_scale, + const WeightValues& w, + uint32_t* out_id) -> TfLiteStatus { + if (!is_quantized) { + TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, /*num_k_dims=*/1, a_id, + w.value_id, YNN_INVALID_VALUE_ID, + out_id, 0)); + return kTfLiteOk; + } + + uint32_t dot_zp = YNN_INVALID_VALUE_ID; + uint32_t dot_scale = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS(ynn::define_dot_quantization( + subgraph, /*num_k_dims=*/1, a_id, a_zp, a_scale, w.value_id, + w.zero_point_id, w.scale_id, dot_zp, dot_scale)); + + // The dot accumulates into the negated zero point correction, so that the + // raw accumulator is already zero-point corrected. + uint32_t accum_init = YNN_INVALID_VALUE_ID; + if (dot_zp != YNN_INVALID_VALUE_ID) { + TF_LITE_ENSURE_YNN_STATUS( + ynn_define_unary(subgraph, ynn_unary_negate, dot_zp, &accum_init, 0)); + } + + uint32_t accum = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot( + subgraph, /*num_k_dims=*/1, a_id, w.value_id, accum_init, &accum, 0)); + + TF_LITE_ENSURE_YNN_STATUS( + ynn_define_dequantize(subgraph, accum, YNN_INVALID_VALUE_ID, dot_scale, + ynn_type_fp32, out_id, 0)); + return kTfLiteOk; + }; + + uint32_t tokens_in = tokens_5d; + uint32_t tokens_zp = YNN_INVALID_VALUE_ID; + uint32_t tokens_scale = YNN_INVALID_VALUE_ID; + if (is_quantized) { + tokens_in = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_STATUS( + dynamic_quantize(tokens_5d, &tokens_in, &tokens_zp, &tokens_scale)); + } // 5. Batched dot products: - // tokens [B, N, 1, 1, D_in] @ w_gate_k [B, N, K, D_in, D_mid] -> gate [B, N, - // K, 1, D_mid] Slinky automatically broadcasts axis 2 (dim 1 -> K) of tokens. uint32_t gate = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, 1, tokens_5d, w_gate_k, - YNN_INVALID_VALUE_ID, &gate, 0)); + TF_LITE_ENSURE_STATUS( + define_dot(tokens_in, tokens_zp, tokens_scale, weights[kGate], &gate)); - uint32_t gelu_gate = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn::define_approx_gelu(subgraph, gate, gelu_gate)); + uint32_t act_gate = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_YNN_STATUS(ynn::define_approx_gelu(subgraph, gate, act_gate)); uint32_t up = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, 1, tokens_5d, w_up_k, - YNN_INVALID_VALUE_ID, &up, 0)); + TF_LITE_ENSURE_STATUS( + define_dot(tokens_in, tokens_zp, tokens_scale, weights[kUp], &up)); uint32_t mid = YNN_INVALID_VALUE_ID; TF_LITE_ENSURE_YNN_STATUS( - ynn_define_binary(subgraph, ynn_binary_multiply, gelu_gate, up, &mid, 0)); + ynn_define_binary(subgraph, ynn_binary_multiply, act_gate, up, &mid, 0)); + + uint32_t mid_in = mid; + uint32_t mid_zp = YNN_INVALID_VALUE_ID; + uint32_t mid_scale = YNN_INVALID_VALUE_ID; + if (is_quantized) { + mid_in = YNN_INVALID_VALUE_ID; + TF_LITE_ENSURE_STATUS(dynamic_quantize(mid, &mid_in, &mid_zp, &mid_scale)); + } else { + // The dots accumulate in fp32, so the intermediate has to be narrowed + // back to match the down weights. + TF_LITE_ENSURE_STATUS(ConvertIfNeeded(context, subgraph, kTfLiteFloat32, + dot_float_type, mid_in)); + } - // mid [B, N, K, 1, D_mid] @ w_down_k [B, N, K, D_mid, D_out] -> out_5d [B, N, - // K, 1, D_out] uint32_t out_5d = YNN_INVALID_VALUE_ID; - TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, 1, mid, w_down_k, - YNN_INVALID_VALUE_ID, &out_5d, 0)); + TF_LITE_ENSURE_STATUS( + define_dot(mid_in, mid_zp, mid_scale, weights[kDown], &out_5d)); // 6. Squeeze axis 3: [B, N, K, 1, D_out] -> [B, N, K, D_out] // TODO: b/558433816 - This should be static_transpose, but it is faster to @@ -242,7 +469,7 @@ TfLiteStatus DefineMoeNode(TfLiteContext* context, ynn_subgraph_t subgraph, TF_LITE_ENSURE_YNN_STATUS(ynn_define_binary(subgraph, ynn_binary_multiply, routing_weights_val, scale_1d, &rw_scaled, 0)); - } else if (num_scale_elements == static_cast(E)) { + } else if (num_scale_elements == static_cast(num_experts)) { uint32_t scale_gathered = YNN_INVALID_VALUE_ID; TF_LITE_ENSURE_YNN_STATUS(ynn_define_gather(subgraph, 1, &gather_axis_0, 3, scale_1d, expert_indices_val, @@ -271,6 +498,14 @@ TfLiteStatus DefineMoeNode(TfLiteContext* context, ynn_subgraph_t subgraph, ynn_define_reduce(subgraph, ynn_reduce_sum, 1, &reduce_axis_2, weighted_out, YNN_INVALID_VALUE_ID, &output_val, 0)); + // The reduction accumulates in fp32; an output value this node created + // itself has to be narrowed to the type of the tensor it stands for. + const TfLiteTensor& output = context->tensors[node.outputs[0]]; + if (out_it == tensor_to_value_id.end()) { + TF_LITE_ENSURE_STATUS(ConvertIfNeeded(context, subgraph, kTfLiteFloat32, + output.type, output_val)); + } + tensor_to_value_id[node.outputs[0]] = output_val; return kTfLiteOk; diff --git a/tensorflow/lite/delegates/ynnpack/moe_bench.cc b/tensorflow/lite/delegates/ynnpack/moe_bench.cc index 5f0e499b14d513..164b0f1731b02e 100644 --- a/tensorflow/lite/delegates/ynnpack/moe_bench.cc +++ b/tensorflow/lite/delegates/ynnpack/moe_bench.cc @@ -14,12 +14,20 @@ limitations under the License. ==============================================================================*/ #include +#include #include #include #include -#include +#include +#include +#include +#include +#include +#include +#include #include +#include "absl/algorithm/container.h" #include "benchmark/benchmark.h" // from @com_google_benchmark #include "flatbuffers/buffer.h" // from @flatbuffers #include "tensorflow/lite/c/c_api_types.h" @@ -33,21 +41,171 @@ namespace tflite { namespace ynnpack { namespace { -// Parameters from Gemma4 26B-A4B MoE model: -// embed_dim = 2816 (D_in / D_out), expert_dim = 704 (D_mid), -// num_experts = 128 (E), top_k = 8 (K). -constexpr int kGemma4Din = 2816; -constexpr int kGemma4Dmid = 704; -constexpr int kGemma4NumExperts = 128; -constexpr int kGemma4TopK = 8; +// Shapes of one MoE block in Qwen1.5-MoE-A2.7B, from the `qwen2_moe` config in +// //third_party/py/transformers (whose defaults are that model): hidden_size = +// 2048, moe_intermediate_size = 1408, num_experts = 60, num_experts_per_tok = +// 4. This is the "large expert" regime: one expert's weights do not fit in L2. +// +// The model also has a shared expert (shared_expert_intermediate_size = 5632) +// that every token passes through. That is a plain dense FFN, not part of the +// routed MoE op measured here, so it is not modelled. +constexpr int kQwenDin = 2048; +constexpr int kQwenDmid = 1408; +constexpr int kQwenNumExperts = 60; +constexpr int kQwenTopK = 4; + +// Shapes of one MoE block in Gemma 4 700M (the "tiny MoE" / E2B model) as used +// for the end-to-end LiteRT LM measurements. These were read back from the +// YNNPACK subgraph dump of the real model, not from a config file. This is the +// "small expert" regime: one expert's weights fit comfortably in L2. +constexpr int kGemma700MDin = 512; +constexpr int kGemma700MDmid = 448; +constexpr int kGemma700MNumExperts = 32; +constexpr int kGemma700MTopK = 4; + +// Shapes of one routed MoE block in Gemma 4 26B-A4B, from the +// `gemma4_26b_a4b` config in //third_party/py/tunix/models/gemma4/config.py +// (cross-checked against `ConfigGemma4_26B_MoE` in +// //third_party/gemma_cpp/gemma/configs.cc and the HF `moe_intermediate_size` +// used in //third_party/py/qwix/google/offline_quantization_example/gemma4.py): +// embed_dim = 2816, expert_dim = 704, num_experts = 128, +// num_experts_per_tok = 8. +// +// Every layer also runs a dense shared MLP (hidden_dim = 2112) in parallel +// with the routed experts. That is a plain FFN, not part of the routed MoE op +// measured here, so it is not modelled. +// +// This is the "many small experts" regime: one expert is ~5.9 MB at int8, +// a third smaller than a Qwen expert, but 8 of the 128 are active per token, +// so a decode step streams ~48 MB of int8 weights against Qwen's ~35 MB and +// 700M's ~2.8 MB, out of a ~760 MB per-block working set. +// +// These arms are expensive: a 26B prefill iteration at seq_len 512 takes +// ~1.4 s at int8, and the FP32 arms allocate ~3 GB of weights, so pass +// --benchmark_filter rather than running the whole file. +constexpr int kGemma26BDin = 2816; +constexpr int kGemma26BDmid = 704; +constexpr int kGemma26BNumExperts = 128; +constexpr int kGemma26BTopK = 8; + +// Number of distinct routing buffers the timed loop cycles through. With a +// single fixed buffer the same few experts stay resident in cache across all +// iterations, which is not what a real decode loop does: there, consecutive +// tokens pick different experts and the weights are streamed from memory. +constexpr int kRoutingRingSize = 64; + +// Generates expert assignments that look like the output of a trained MoE +// router rather than a round-robin counter. +// +// Experts are given Zipf-distributed popularities w[e] ~ 1 / rank[e]^skew, +// where `rank` is a random permutation of the expert ids so that popular +// experts are not adjacent in memory. Each token then draws K *distinct* +// experts with probability proportional to w, using the Gumbel top-k trick +// (equivalent to Plackett-Luce sampling without replacement). +// +// The `skew` knob is not a load imbalance in itself; the resulting max/mean +// expert load also depends on E and K. Measured over 512 tokens: +// +// skew E=32,K=4 E=60,K=4 +// 0.0 1.27 1.40 (floor: pure sampling noise) +// 0.2 1.58 1.88 +// 0.3 1.89 2.44 +// 0.7 3.91 6.05 +// 1.2 6.63 11.78 +// +// A model trained with a load-balancing auxiliary loss usually lands near +// 1.2-2.0, i.e. skew in the 0.2-0.3 range. Rather than trusting the knob, +// read the `load_imbalance` counter the benchmark emits. +class ExpertRouter { + public: + ExpertRouter(int num_experts, int top_k, double skew, uint64_t seed) + : num_experts_(num_experts), top_k_(top_k), rng_(seed) { + std::vector rank(num_experts); + std::iota(rank.begin(), rank.end(), 0); + std::shuffle(rank.begin(), rank.end(), rng_); + log_weight_.resize(num_experts); + for (int e = 0; e < num_experts; ++e) { + log_weight_[e] = -skew * std::log(static_cast(rank[e]) + 1.0); + } + } + + // Fills `out` with `num_tokens * top_k` expert ids. + void Route(int num_tokens, int32_t* out) { + std::uniform_real_distribution uniform( + std::numeric_limits::min(), 1.0); + std::vector> keys(num_experts_); + for (int t = 0; t < num_tokens; ++t) { + for (int e = 0; e < num_experts_; ++e) { + const double gumbel = -std::log(-std::log(uniform(rng_))); + keys[e] = {log_weight_[e] + gumbel, e}; + } + absl::c_partial_sort( + keys, keys.begin() + top_k_, + [](const std::pair& a, const std::pair& b) { + return a.first > b.first; + }); + for (int k = 0; k < top_k_; ++k) { + out[t * top_k_ + k] = keys[k].second; + } + } + } + + private: + int num_experts_; + int top_k_; + std::mt19937_64 rng_; + std::vector log_weight_; +}; + +struct RoutingStats { + // Mean over invocations of (max expert load / mean expert load) within one + // invocation. 1.0 is perfectly balanced. + double load_imbalance = 0.0; + // Mean over invocations of the number of distinct experts one invocation + // touches. + double experts_per_invocation = 0.0; + // Number of distinct experts touched across the whole routing ring. Times + // the per-expert weight footprint, this is the working set the timed loop + // actually streams. + int experts_in_ring = 0; +}; + +RoutingStats ComputeRoutingStats( + const std::vector>& routing, int num_experts) { + RoutingStats stats; + std::vector ring_seen(num_experts, 0); + for (const std::vector& buffer : routing) { + std::vector count(num_experts, 0); + for (int32_t e : buffer) { + ++count[e]; + ring_seen[e] = 1; + } + int max_count = 0; + int distinct = 0; + for (int c : count) { + max_count = std::max(max_count, c); + distinct += c > 0 ? 1 : 0; + } + const double mean = + static_cast(buffer.size()) / static_cast(num_experts); + stats.load_imbalance += static_cast(max_count) / mean; + stats.experts_per_invocation += distinct; + } + const double n = static_cast(routing.size()); + stats.load_imbalance /= n; + stats.experts_per_invocation /= n; + stats.experts_in_ring = + std::accumulate(ring_seen.begin(), ring_seen.end(), 0); + return stats; +} class MoeModel : public SingleOpModel { public: MoeModel(int B, int N, int D_in, int D_mid, int E, int K, int num_threads, bool use_delegate) - : w_gate_data_(D_mid * E * D_in), - w_up_data_(D_mid * E * D_in), - w_down_data_(D_in * E * D_mid), + : w_gate_data_(static_cast(D_mid) * E * D_in), + w_up_data_(static_cast(D_mid) * E * D_in), + w_down_data_(static_cast(D_in) * E * D_mid), scale_data_(E) { std::vector tokens_shape = {B, N, D_in}; std::vector rw_shape = {B, N, K}; @@ -106,6 +264,10 @@ class MoeModel : public SingleOpModel { reinterpret_cast(scale_data_.data()), scale_data_.size() * sizeof(float)); + // Note: there is no undelegated arm here. The model contains a single + // stablehlo.composite op whose decomposition subgraph is not materialized + // by SingleOpModel, so it can only run through the YNNPACK delegate. Use + // the end-to-end LiteRT LM benchmark to get a delegated/undelegated ratio. if (use_delegate) { TfLiteYNNPackDelegateOptions delegate_options = TfLiteYNNPackDelegateOptionsDefault(); @@ -126,6 +288,9 @@ class MoeModel : public SingleOpModel { int ei() const { return ei_id_; } int out() const { return out_id_; } + // Size of one stored weight element, for the working-set counter. + static constexpr double kWeightBytes = sizeof(float); + private: int tokens_id_; int rw_id_; @@ -142,83 +307,569 @@ class MoeModel : public SingleOpModel { std::vector scale_data_; }; -void BenchMoe(benchmark::State& state, int B, int N, int D_in, int D_mid, int E, - int K, int num_threads) { - MoeModel model(B, N, D_in, D_mid, E, K, num_threads, /*use_delegate=*/true); - - std::vector tokens(B * N * D_in, 0.5f); - std::vector rw(B * N * K, 1.0f / static_cast(K)); - std::vector ei(B * N * K); - for (int n = 0; n < B * N; ++n) { - for (int k = 0; k < K; ++k) { - ei[n * K + k] = (n * K + k) % E; +// Storage precision of the expert weights in `QuantizedMoeModel`. int4 and +// int2 weights are stored packed, 2 respectively 4 values per byte. +enum class WeightPrecision { kInt8, kInt4, kInt2 }; + +template +struct WeightTraits; + +template <> +struct WeightTraits { + static constexpr TensorType kTensorType = TensorType_INT8; + static constexpr TfLiteType kTfLiteType = kTfLiteInt8; + static constexpr int kValuesPerByte = 1; + // Weights cycle through [-kValuePeriod / 2, kValuePeriod / 2]. + static constexpr int kValuePeriod = 31; +}; + +template <> +struct WeightTraits { + static constexpr TensorType kTensorType = TensorType_INT4; + static constexpr TfLiteType kTfLiteType = kTfLiteInt4; + static constexpr int kValuesPerByte = 2; + static constexpr int kValuePeriod = 15; +}; + +template <> +struct WeightTraits { + static constexpr TensorType kTensorType = TensorType_INT2; + static constexpr TfLiteType kTfLiteType = kTfLiteInt2; + static constexpr int kValuesPerByte = 4; + static constexpr int kValuePeriod = 3; +}; + +// The dynamically quantized counterpart of `MoeModel`, with the weights stored +// at the precision `P` (W8A8, W4A8 or W2A8). +// +// This is the 10-input form of the op: each weight tensor is followed by its +// own FP32 per-output-channel scale tensor. Activations are quantized by the +// delegate at run time, so the inputs and the output stay FP32 and the +// benchmark loop is identical to the float one. +// +// Like `MoeModel`, the weights are attached with +// `SetTensorParametersReadOnly` rather than `AddConstInput`. At Qwen shapes +// the three weight tensors are ~600 MB and embedding them in the flatbuffer +// would be slow and would risk the 2 GB flatbuffer limit. It also matters +// for what this benchmark measures: read-only tensors are what the delegate +// treats as constant, which is the precondition for the weight-gather +// fusions to fire. +template +class QuantizedMoeModel : public SingleOpModel { + public: + QuantizedMoeModel(int B, int N, int D_in, int D_mid, int E, int K, + int num_threads, bool use_delegate) + : w_gate_data_(static_cast(D_mid) * E * D_in), + w_up_data_(static_cast(D_mid) * E * D_in), + w_down_data_(static_cast(D_in) * E * D_mid), + gate_scale_data_(static_cast(D_mid) * E), + up_scale_data_(static_cast(D_mid) * E), + down_scale_data_(static_cast(D_in) * E), + scale_data_(E) { + std::vector tokens_shape = {B, N, D_in}; + std::vector rw_shape = {B, N, K}; + std::vector ei_shape = {B, N, K}; + std::vector w_gate_shape = {D_mid, E, 1, D_in}; + std::vector gate_scale_shape = {D_mid, E, 1, 1}; + std::vector w_up_shape = {D_mid, E, 1, D_in}; + std::vector up_scale_shape = {D_mid, E, 1, 1}; + std::vector w_down_shape = {D_in, E, 1, D_mid}; + std::vector down_scale_shape = {D_in, E, 1, 1}; + std::vector scale_shape = {1, 1, 1, E}; + std::vector out_shape = {B, N, D_in}; + + tokens_id_ = AddInput({TensorType_FLOAT32, tokens_shape}); + rw_id_ = AddInput({TensorType_FLOAT32, rw_shape}); + ei_id_ = AddInput({TensorType_INT32, ei_shape}); + w_gate_id_ = AddInput({WeightTraits

::kTensorType, w_gate_shape}); + gate_scale_id_ = AddInput({TensorType_FLOAT32, gate_scale_shape}); + w_up_id_ = AddInput({WeightTraits

::kTensorType, w_up_shape}); + up_scale_id_ = AddInput({TensorType_FLOAT32, up_scale_shape}); + w_down_id_ = AddInput({WeightTraits

::kTensorType, w_down_shape}); + down_scale_id_ = AddInput({TensorType_FLOAT32, down_scale_shape}); + scale_id_ = AddInput({TensorType_FLOAT32, scale_shape}); + + out_id_ = AddOutput({TensorType_FLOAT32, out_shape}); + + std::vector empty_attrs; + flatbuffers::Offset options = + CreateStableHLOCompositeOptionsDirect( + builder_, "odml.moe_experts", + /*decomposition_subgraph_index=*/1, &empty_attrs); + + SetBuiltinOp(BuiltinOperator_STABLEHLO_COMPOSITE, + BuiltinOptions2_StableHLOCompositeOptions, options.Union()); + + BuildInterpreter({tokens_shape, rw_shape, ei_shape, w_gate_shape, + gate_scale_shape, w_up_shape, up_scale_shape, + w_down_shape, down_scale_shape, scale_shape}, + -1, false, false, + /*allocate_and_delegate=*/false); + + // Spread the weights over the representable range rather than filling with + // a constant, so that neither the quantized dot nor the zero-point + // correction can be short-circuited by a degenerate value. + constexpr int period = WeightTraits

::kValuePeriod; + for (size_t i = 0; i < w_gate_data_.size(); ++i) { + w_gate_data_[i] = static_cast((i % period) - period / 2); + } + for (size_t i = 0; i < w_up_data_.size(); ++i) { + w_up_data_[i] = static_cast(((i + 7) % period) - period / 2); + } + for (size_t i = 0; i < w_down_data_.size(); ++i) { + w_down_data_[i] = static_cast(((i + 13) % period) - period / 2); + } + std::fill(gate_scale_data_.begin(), gate_scale_data_.end(), 0.005f); + std::fill(up_scale_data_.begin(), up_scale_data_.end(), 0.005f); + std::fill(down_scale_data_.begin(), down_scale_data_.end(), 0.005f); + std::fill(scale_data_.begin(), scale_data_.end(), 1.0f); + + SetReadOnlyWeights(w_gate_id_, "gate_weights", w_gate_shape, w_gate_data_); + SetReadOnlyWeights(w_up_id_, "up_weights", w_up_shape, w_up_data_); + SetReadOnlyWeights(w_down_id_, "down_weights", w_down_shape, w_down_data_); + SetReadOnlyFloat(gate_scale_id_, "gate_scale", gate_scale_shape, + gate_scale_data_); + SetReadOnlyFloat(up_scale_id_, "up_scale", up_scale_shape, up_scale_data_); + SetReadOnlyFloat(down_scale_id_, "down_scale", down_scale_shape, + down_scale_data_); + SetReadOnlyFloat(scale_id_, "scale", scale_shape, scale_data_); + + if (use_delegate) { + TfLiteYNNPackDelegateOptions delegate_options = + TfLiteYNNPackDelegateOptionsDefault(); + delegate_options.num_threads = num_threads; + delegate_options.static_shape = true; + SetDelegate(Interpreter::TfLiteDelegatePtr( + TfLiteYNNPackDelegateCreate(&delegate_options), + TfLiteYNNPackDelegateDelete)); + ApplyDelegate(); + } + if (interpreter_->AllocateTensors() != kTfLiteOk) { + fprintf(stderr, "Failed to allocate tensors\n"); + } + } + + int tokens() const { return tokens_id_; } + int rw() const { return rw_id_; } + int ei() const { return ei_id_; } + int out() const { return out_id_; } + + static constexpr double kWeightBytes = 1.0 / WeightTraits

::kValuesPerByte; + + private: + // Packs `data` the way LiteRT stores sub-byte tensors (lowest-index value in + // the least significant bits) and hands the result to the interpreter as a + // read-only tensor. `packed_weights_` keeps the storage alive: the + // interpreter does not copy it. + void SetReadOnlyWeights(int id, const char* name, + const std::vector& shape, + const std::vector& data) { + constexpr int kValuesPerByte = WeightTraits

::kValuesPerByte; + const char* bytes = reinterpret_cast(data.data()); + size_t num_bytes = data.size(); + if (kValuesPerByte > 1) { + constexpr int kBits = 8 / kValuesPerByte; + constexpr int kMask = (1 << kBits) - 1; + std::vector packed( + (data.size() + kValuesPerByte - 1) / kValuesPerByte, 0); + for (size_t i = 0; i < data.size(); ++i) { + packed[i / kValuesPerByte] |= static_cast( + (data[i] & kMask) << (kBits * (i % kValuesPerByte))); + } + num_bytes = packed.size(); + packed_weights_.push_back(std::move(packed)); + bytes = reinterpret_cast(packed_weights_.back().data()); } + interpreter_->SetTensorParametersReadOnly(id, WeightTraits

::kTfLiteType, + name, shape, TfLiteQuantization(), + bytes, num_bytes); + } + + void SetReadOnlyFloat(int id, const char* name, const std::vector& shape, + const std::vector& data) { + interpreter_->SetTensorParametersReadOnly( + id, kTfLiteFloat32, name, shape, TfLiteQuantization(), + reinterpret_cast(data.data()), + data.size() * sizeof(float)); + } + + int tokens_id_; + int rw_id_; + int ei_id_; + int w_gate_id_; + int gate_scale_id_; + int w_up_id_; + int up_scale_id_; + int w_down_id_; + int down_scale_id_; + int scale_id_; + int out_id_; + + std::vector w_gate_data_; + std::vector w_up_data_; + std::vector w_down_data_; + // Packed copies of the weights above, for the sub-byte precisions. + std::vector> packed_weights_; + std::vector gate_scale_data_; + std::vector up_scale_data_; + std::vector down_scale_data_; + std::vector scale_data_; +}; + +// `SingleOpModel::ApplyDelegate` logs a warning every time a delegate is set +// manually, and `TFLITE_LOG` writes to `std::cout` unconditionally: it has no +// severity filter, and despite its comment it does not use stderr. Google +// Benchmark also writes its report to `std::cout`, so with +// `--benchmark_format=json` those lines land in the middle of the JSON +// document and make it unparseable. benchy then fails the whole run with +// "invalid character 'W' looking for beginning of value". +// +// Point `std::cout` at `std::cerr` for the duration of the noisy window so +// the diagnostics are still visible, just not on the machine-readable +// stream. +class ScopedCoutToCerr { + public: + ScopedCoutToCerr() : saved_(std::cout.rdbuf(std::cerr.rdbuf())) {} + ~ScopedCoutToCerr() { std::cout.rdbuf(saved_); } + + ScopedCoutToCerr(const ScopedCoutToCerr&) = delete; + ScopedCoutToCerr& operator=(const ScopedCoutToCerr&) = delete; + + private: + std::streambuf* saved_; +}; + +template +void BenchMoe(benchmark::State& state, int B, int N, int D_in, int D_mid, int E, + int K, int num_threads, double skew, bool vary_routing) { + // Held for the whole body: the benchmark framework emits its report after + // this function returns, so nothing of ours is diverted. + ScopedCoutToCerr redirect_tflite_logs; + + ModelT model(B, N, D_in, D_mid, E, K, num_threads, /*use_delegate=*/true); + + const int num_tokens = B * N; + std::vector tokens(static_cast(num_tokens) * D_in, 0.5f); + std::vector rw(static_cast(num_tokens) * K, + 1.0f / static_cast(K)); + + const int ring_size = vary_routing ? kRoutingRingSize : 1; + ExpertRouter router(E, K, skew, /*seed=*/0x5eed1234u); + std::vector> routing(ring_size); + for (std::vector& buffer : routing) { + buffer.resize(static_cast(num_tokens) * K); + router.Route(num_tokens, buffer.data()); } model.PopulateTensor(model.tokens(), tokens); model.PopulateTensor(model.rw(), rw); - model.PopulateTensor(model.ei(), ei); + model.PopulateTensor(model.ei(), routing[0]); if (model.Invoke() != kTfLiteOk) { state.SkipWithError("Failed to invoke delegate interpreter (warmup)"); return; } + // The routing tensor is repopulated on every iteration in both arms so that + // the fixed and varying arms differ only in the *content* of the routing, + // never in the work done outside the delegate. + int index = 0; for (auto _ : state) { + model.PopulateTensor(model.ei(), routing[index]); + index = index + 1 == ring_size ? 0 : index + 1; if (model.Invoke() != kTfLiteOk) { state.SkipWithError("Failed to invoke interpreter"); return; } } - // 2 matmuls of [D_in, D_mid] (gate and up) + 1 matmul of [D_mid, D_in] (down) - // per expert choice. - const size_t flops = 6ull * static_cast(B) * static_cast(N) * + // 2 matmuls of [D_in, D_mid] (gate and up) + 1 matmul of [D_mid, D_in] + // (down) per (token, expert) pair. + const size_t flops = 6ull * static_cast(num_tokens) * static_cast(K) * static_cast(D_in) * static_cast(D_mid); - state.counters["FLOP"] = - benchmark::Counter(static_cast(state.iterations() * flops), - benchmark::Counter::kIsRate); + state.counters["FLOP"] = benchmark::Counter( + static_cast(state.iterations()) * static_cast(flops), + benchmark::Counter::kIsRate); + + const RoutingStats stats = ComputeRoutingStats(routing, E); + state.counters["load_imbalance"] = stats.load_imbalance; + state.counters["experts_per_invoke"] = stats.experts_per_invocation; + // Weight bytes touched over one full cycle of the routing ring. + const double expert_bytes = + 3.0 * static_cast(D_in) * D_mid * ModelT::kWeightBytes; + state.counters["working_set_MB"] = + stats.experts_in_ring * expert_bytes / (1024.0 * 1024.0); +} + +// The argument unpacking is shared so that the FP32 and INT8 arms of a given +// shape regime are guaranteed to see identical sequence lengths, thread +// counts, routing skew and routing ring. +template +void DecodeBench(benchmark::State& state, int D_in, int D_mid, int E, int K) { + BenchMoe(state, /*B=*/1, /*N=*/1, D_in, D_mid, E, K, + /*num_threads=*/state.range(0), + /*skew=*/state.range(1) / 10.0, + /*vary_routing=*/state.range(2) != 0); +} + +template +void PrefillBench(benchmark::State& state, int D_in, int D_mid, int E, int K) { + BenchMoe(state, /*B=*/1, /*N=*/state.range(0), D_in, D_mid, E, K, + /*num_threads=*/state.range(1), + /*skew=*/state.range(2) / 10.0, /*vary_routing=*/true); +} + +void QwenMoeDecodeFloat32(benchmark::State& state) { + DecodeBench(state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void QwenMoeDecodeInt8(benchmark::State& state) { + DecodeBench>( + state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void QwenMoeDecodeInt4(benchmark::State& state) { + DecodeBench>( + state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void QwenMoeDecodeInt2(benchmark::State& state) { + DecodeBench>( + state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void QwenMoePrefillFloat32(benchmark::State& state) { + PrefillBench(state, kQwenDin, kQwenDmid, kQwenNumExperts, + kQwenTopK); +} + +void QwenMoePrefillInt8(benchmark::State& state) { + PrefillBench>( + state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void QwenMoePrefillInt4(benchmark::State& state) { + PrefillBench>( + state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void QwenMoePrefillInt2(benchmark::State& state) { + PrefillBench>( + state, kQwenDin, kQwenDmid, kQwenNumExperts, kQwenTopK); +} + +void Gemma700MMoeDecodeFloat32(benchmark::State& state) { + DecodeBench(state, kGemma700MDin, kGemma700MDmid, + kGemma700MNumExperts, kGemma700MTopK); +} + +void Gemma700MMoeDecodeInt8(benchmark::State& state) { + DecodeBench>( + state, kGemma700MDin, kGemma700MDmid, kGemma700MNumExperts, + kGemma700MTopK); +} + +void Gemma700MMoeDecodeInt4(benchmark::State& state) { + DecodeBench>( + state, kGemma700MDin, kGemma700MDmid, kGemma700MNumExperts, + kGemma700MTopK); +} + +void Gemma700MMoeDecodeInt2(benchmark::State& state) { + DecodeBench>( + state, kGemma700MDin, kGemma700MDmid, kGemma700MNumExperts, + kGemma700MTopK); +} + +void Gemma700MMoePrefillFloat32(benchmark::State& state) { + PrefillBench(state, kGemma700MDin, kGemma700MDmid, + kGemma700MNumExperts, kGemma700MTopK); +} + +void Gemma700MMoePrefillInt8(benchmark::State& state) { + PrefillBench>( + state, kGemma700MDin, kGemma700MDmid, kGemma700MNumExperts, + kGemma700MTopK); +} + +void Gemma700MMoePrefillInt4(benchmark::State& state) { + PrefillBench>( + state, kGemma700MDin, kGemma700MDmid, kGemma700MNumExperts, + kGemma700MTopK); +} + +void Gemma700MMoePrefillInt2(benchmark::State& state) { + PrefillBench>( + state, kGemma700MDin, kGemma700MDmid, kGemma700MNumExperts, + kGemma700MTopK); +} + +// 26B-A4B holds 3 * 2816 * 704 * 128 = 761M weight values per MoE block, so +// the three weight tensors are ~3.0 GB in FP32, ~760 MB at int8, ~380 MB at +// int4 and ~190 MB at int2. The FP32 arm is a reference point, not a +// deployable configuration. +void Gemma26BMoeDecodeFloat32(benchmark::State& state) { + DecodeBench(state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, + kGemma26BTopK); +} + +void Gemma26BMoeDecodeInt8(benchmark::State& state) { + DecodeBench>( + state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, kGemma26BTopK); +} + +void Gemma26BMoeDecodeInt4(benchmark::State& state) { + DecodeBench>( + state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, kGemma26BTopK); +} + +void Gemma26BMoeDecodeInt2(benchmark::State& state) { + DecodeBench>( + state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, kGemma26BTopK); +} + +void Gemma26BMoePrefillFloat32(benchmark::State& state) { + PrefillBench(state, kGemma26BDin, kGemma26BDmid, + kGemma26BNumExperts, kGemma26BTopK); +} + +void Gemma26BMoePrefillInt8(benchmark::State& state) { + PrefillBench>( + state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, kGemma26BTopK); } -void Gemma4MoEDecode(benchmark::State& state) { - const int num_threads = state.range(0); - BenchMoe(state, /*B=*/1, /*N=*/1, kGemma4Din, kGemma4Dmid, kGemma4NumExperts, - kGemma4TopK, num_threads); +void Gemma26BMoePrefillInt4(benchmark::State& state) { + PrefillBench>( + state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, kGemma26BTopK); } -void Gemma4MoEPrefill(benchmark::State& state) { - const int seq_len = state.range(0); - const int num_threads = state.range(1); - BenchMoe(state, /*B=*/1, seq_len, kGemma4Din, kGemma4Dmid, kGemma4NumExperts, - kGemma4TopK, num_threads); +void Gemma26BMoePrefillInt2(benchmark::State& state) { + PrefillBench>( + state, kGemma26BDin, kGemma26BDmid, kGemma26BNumExperts, kGemma26BTopK); } void DecodeArguments(benchmark::internal::Benchmark* b) { - b->ArgNames({"threads"}); + b->ArgNames({"threads", "skew_x10", "vary_routing"}); b->UseRealTime(); b->MeasureProcessCPUTime(); for (int threads : {1, 4}) { - b->Args({threads}); + // vary_routing = 0 keeps the same experts hot in cache across iterations + // and is only kept so the cache effect can be quantified; 1 is the + // realistic arm. + for (int vary : {0, 1}) { + b->Args({threads, 3, vary}); + } } } void PrefillArguments(benchmark::internal::Benchmark* b) { - b->ArgNames({"seq_len", "threads"}); + b->ArgNames({"seq_len", "threads", "skew_x10"}); b->UseRealTime(); b->MeasureProcessCPUTime(); - for (int seq_len : {64, 256, 1024}) { + for (int seq_len : {128, 256, 512}) { for (int threads : {1, 4}) { - b->Args({seq_len, threads}); + // Uniform / realistic / pathological expert load. + for (int skew : {0, 3, 12}) { + b->Args({seq_len, threads, skew}); + } } } } -BENCHMARK(Gemma4MoEDecode) +BENCHMARK(QwenMoeDecodeFloat32) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoeDecodeInt8) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoeDecodeInt4) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoeDecodeInt2) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoePrefillFloat32) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoePrefillInt8) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoePrefillInt4) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(QwenMoePrefillInt2) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoeDecodeFloat32) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoeDecodeInt8) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoeDecodeInt4) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoeDecodeInt2) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoePrefillFloat32) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoePrefillInt8) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoePrefillInt4) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma700MMoePrefillInt2) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoeDecodeFloat32) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoeDecodeInt8) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoeDecodeInt4) ->Apply(DecodeArguments) ->Unit(benchmark::TimeUnit::kMillisecond); -BENCHMARK(Gemma4MoEPrefill) +BENCHMARK(Gemma26BMoeDecodeInt2) + ->Apply(DecodeArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoePrefillFloat32) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoePrefillInt8) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoePrefillInt4) + ->Apply(PrefillArguments) + ->Unit(benchmark::TimeUnit::kMillisecond); + +BENCHMARK(Gemma26BMoePrefillInt2) ->Apply(PrefillArguments) ->Unit(benchmark::TimeUnit::kMillisecond); diff --git a/tensorflow/lite/delegates/ynnpack/moe_test.cc b/tensorflow/lite/delegates/ynnpack/moe_test.cc index 234056ee76a1e8..788ddb25ffc16d 100644 --- a/tensorflow/lite/delegates/ynnpack/moe_test.cc +++ b/tensorflow/lite/delegates/ynnpack/moe_test.cc @@ -19,14 +19,17 @@ limitations under the License. #include #include #include +#include #include #include +#include "Eigen/Core" // from @eigen_archive #include "flatbuffers/buffer.h" // from @flatbuffers #include "tensorflow/lite/delegates/ynnpack/ynnpack_delegate.h" #include "tensorflow/lite/interpreter.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" +#include "tensorflow/lite/types/half.h" namespace tflite { namespace ynnpack { @@ -169,8 +172,8 @@ void RunReferenceMoe(int B, int N, int D_in, int D_mid, int E, int K, sum_u += token[d_in] * w_up[src_idx]; } float x = sum_g; - float gelu = x * 0.5f * (1.0f + std::erf(x * 0.70710678118654752440f)); - gate_val[d_mid] = gelu * sum_u; + float act = x * 0.5f * (1.0f + std::erf(x * 0.70710678118654752440f)); + gate_val[d_mid] = act * sum_u; } for (int d_in = 0; d_in < D_in; ++d_in) { @@ -526,6 +529,687 @@ TEST(MoeTest, ResizeTensorDefaultOptions) { } } +template +class FloatMoeModel : public SingleOpModel { + public: + FloatMoeModel(int B, int N, int D_in, int D_mid, int E, int K, + bool use_delegate) + : w_gate_ref_(D_mid * E * D_in), + w_up_ref_(D_mid * E * D_in), + w_down_ref_(D_in * E * D_mid), + scale_ref_(E) { + std::vector tokens_shape = {B, N, D_in}; + std::vector rw_shape = {B, N, K}; + std::vector ei_shape = {B, N, K}; + std::vector w_gate_shape = {D_mid, E, 1, D_in}; + std::vector w_up_shape = {D_mid, E, 1, D_in}; + std::vector w_down_shape = {D_in, E, 1, D_mid}; + std::vector scale_shape = {1, 1, 1, E}; + std::vector out_shape = {B, N, D_in}; + + tokens_id_ = AddInput({GetTensorType(), tokens_shape}); + rw_id_ = AddInput({GetTensorType(), rw_shape}); + ei_id_ = AddInput({TensorType_INT32, ei_shape}); + + std::vector w_gate_data(w_gate_ref_.size()); + std::vector w_up_data(w_up_ref_.size()); + std::vector w_down_data(w_down_ref_.size()); + std::vector scale_data(scale_ref_.size()); + + for (size_t i = 0; i < w_gate_data.size(); ++i) { + float v = 0.1f * std::sin(static_cast(i)); + w_gate_data[i] = static_cast(v); + w_gate_ref_[i] = static_cast(w_gate_data[i]); + } + for (size_t i = 0; i < w_up_data.size(); ++i) { + float v = 0.1f * std::sin(static_cast(i) + 1.0f); + w_up_data[i] = static_cast(v); + w_up_ref_[i] = static_cast(w_up_data[i]); + } + for (size_t i = 0; i < w_down_data.size(); ++i) { + float v = 0.1f * std::sin(static_cast(i) + 2.0f); + w_down_data[i] = static_cast(v); + w_down_ref_[i] = static_cast(w_down_data[i]); + } + for (size_t i = 0; i < scale_data.size(); ++i) { + float v = 1.0f + 0.1f * std::sin(static_cast(i) + 3.0f); + scale_data[i] = static_cast(v); + scale_ref_[i] = static_cast(scale_data[i]); + } + + w_gate_id_ = AddConstInput( + TensorData{GetTensorType(), w_gate_shape}, w_gate_data); + w_up_id_ = AddConstInput(TensorData{GetTensorType(), w_up_shape}, + w_up_data); + w_down_id_ = AddConstInput( + TensorData{GetTensorType(), w_down_shape}, w_down_data); + scale_id_ = AddConstInput(TensorData{GetTensorType(), scale_shape}, + scale_data); + + out_id_ = AddOutput({GetTensorType(), out_shape}); + + std::vector empty_attrs; + flatbuffers::Offset options = + CreateStableHLOCompositeOptionsDirect( + builder_, "odml.moe_experts", + /*decomposition_subgraph_index=*/1, &empty_attrs); + + SetBuiltinOp(BuiltinOperator_STABLEHLO_COMPOSITE, + BuiltinOptions2_StableHLOCompositeOptions, options.Union()); + + BuildInterpreter({tokens_shape, rw_shape, ei_shape}, -1, false, false, + /*allocate_and_delegate=*/false); + + if (use_delegate) { + TfLiteYNNPackDelegateOptions delegate_options = + TfLiteYNNPackDelegateOptionsDefault(); + delegate_options.num_threads = 1; + delegate_options.static_shape = true; + SetDelegate(Interpreter::TfLiteDelegatePtr( + TfLiteYNNPackDelegateCreate(&delegate_options), + TfLiteYNNPackDelegateDelete)); + ApplyDelegate(); + } + if (interpreter_->AllocateTensors() != kTfLiteOk) { + fprintf(stderr, "Failed to allocate tensors\n"); + } + } + + int tokens() const { return tokens_id_; } + int rw() const { return rw_id_; } + int ei() const { return ei_id_; } + int out() const { return out_id_; } + + const std::vector& w_gate_ref() const { return w_gate_ref_; } + const std::vector& w_up_ref() const { return w_up_ref_; } + const std::vector& w_down_ref() const { return w_down_ref_; } + const std::vector& scale_ref() const { return scale_ref_; } + + private: + int tokens_id_; + int rw_id_; + int ei_id_; + int w_gate_id_; + int w_up_id_; + int w_down_id_; + int scale_id_; + int out_id_; + + std::vector w_gate_ref_; + std::vector w_up_ref_; + std::vector w_down_ref_; + std::vector scale_ref_; +}; + +// Number of values LiteRT packs into a single byte for `type`. +int ValuesPerByte(TensorType type) { + switch (type) { + case TensorType_INT4: + return 2; + case TensorType_INT2: + return 4; + default: + return 1; + } +} + +// Pseudo-random quantized weight values covering the representable range of +// `type`: [-15, 15] for int8 (the exact range does not matter, it just has to +// be symmetric), [-7, 7] for int4 as the LiteRT quantization spec prescribes, +// and [-2, 1] for int2, which is a plain two's complement 2-bit value. +// +// The values must not be a short cyclic pattern: a period that divides the +// contraction dimension makes every row of the weight matrix identical, and +// the reference output then collapses to near zero and compares equal to +// anything, including a delegate that returns all zeros. +std::vector MakeQuantizedWeights(size_t count, uint32_t seed, + TensorType type) { + int min_value = -15; + int max_value = 15; + switch (type) { + case TensorType_INT4: + min_value = -7; + max_value = 7; + break; + case TensorType_INT2: + min_value = -2; + max_value = 1; + break; + default: + break; + } + const uint32_t range = max_value - min_value + 1; + std::vector values(count); + uint32_t state = seed; + for (size_t i = 0; i < count; ++i) { + state = state * 1664525u + 1013904223u; + values[i] = static_cast(min_value + + static_cast((state >> 16) % range)); + } + return values; +} + +// Scale of the quantized weights, chosen per type so that the MoE output has a +// comparable magnitude whatever the weight precision is. +float WeightScaleFor(TensorType type) { + switch (type) { + case TensorType_INT4: + return 0.04f; + case TensorType_INT2: + return 0.15f; + default: + return 0.02f; + } +} + +// Packs `values` the way LiteRT stores sub-byte tensors: the lowest-index value +// goes into the least significant bits of the byte. +std::vector PackQuantizedWeights(const std::vector& values, + TensorType type) { + const int values_per_byte = ValuesPerByte(type); + const int bits = 8 / values_per_byte; + const int mask = (1 << bits) - 1; + std::vector packed( + (values.size() + values_per_byte - 1) / values_per_byte, 0); + for (size_t i = 0; i < values.size(); ++i) { + packed[i / values_per_byte] |= static_cast( + (values[i] & mask) << (bits * (i % values_per_byte))); + } + return packed; +} + +// `ActT` is the type of the tokens, the routing weights and the output. +template +class QuantizedMoeModel : public SingleOpModel { + public: + QuantizedMoeModel(int B, int N, int D_in, int D_mid, int E, int K, + bool use_delegate, bool explicit_scales = true, + TensorType scale_type = TensorType_FLOAT32, + TensorType weight_type = TensorType_INT8) + : weight_scale_(WeightScaleFor(weight_type)) { + std::vector tokens_shape = {B, N, D_in}; + std::vector rw_shape = {B, N, K}; + std::vector ei_shape = {B, N, K}; + std::vector w_gate_shape = {D_mid, E, 1, D_in}; + std::vector gate_scale_shape = {D_mid, E, 1, 1}; + std::vector w_up_shape = {D_mid, E, 1, D_in}; + std::vector up_scale_shape = {D_mid, E, 1, 1}; + std::vector w_down_shape = {D_in, E, 1, D_mid}; + std::vector down_scale_shape = {D_in, E, 1, 1}; + std::vector scale_shape = {1, 1, 1, E}; + std::vector out_shape = {B, N, D_in}; + + tokens_id_ = AddInput({GetTensorType(), tokens_shape}); + rw_id_ = AddInput({GetTensorType(), rw_shape}); + ei_id_ = AddInput({TensorType_INT32, ei_shape}); + + const size_t gate_up_size = static_cast(D_mid) * E * D_in; + const size_t down_size = static_cast(D_in) * E * D_mid; + const std::vector w_gate_data = + MakeQuantizedWeights(gate_up_size, /*seed=*/1, weight_type); + const std::vector w_up_data = + MakeQuantizedWeights(gate_up_size, /*seed=*/2, weight_type); + const std::vector w_down_data = + MakeQuantizedWeights(down_size, /*seed=*/3, weight_type); + + w_gate_ref_ = Dequantize(w_gate_data); + w_up_ref_ = Dequantize(w_up_data); + w_down_ref_ = Dequantize(w_down_data); + scale_ref_.assign(E, 1.0f); + + const std::vector w_gate_packed = + PackQuantizedWeights(w_gate_data, weight_type); + const std::vector w_up_packed = + PackQuantizedWeights(w_up_data, weight_type); + const std::vector w_down_packed = + PackQuantizedWeights(w_down_data, weight_type); + + if (explicit_scales) { + w_gate_id_ = + AddConstInput(TensorData{weight_type, w_gate_shape}, w_gate_packed); + gate_scale_id_ = + AddScaleInput(scale_type, gate_scale_shape, gate_up_size / D_in); + w_up_id_ = + AddConstInput(TensorData{weight_type, w_up_shape}, w_up_packed); + up_scale_id_ = + AddScaleInput(scale_type, up_scale_shape, gate_up_size / D_in); + w_down_id_ = + AddConstInput(TensorData{weight_type, w_down_shape}, w_down_packed); + down_scale_id_ = + AddScaleInput(scale_type, down_scale_shape, down_size / D_mid); + } else { + w_gate_id_ = AddConstInput( + TensorData{weight_type, w_gate_shape, 0.0f, 0.0f, weight_scale_, 0}, + w_gate_packed); + w_up_id_ = AddConstInput( + TensorData{weight_type, w_up_shape, 0.0f, 0.0f, weight_scale_, 0}, + w_up_packed); + w_down_id_ = AddConstInput( + TensorData{weight_type, w_down_shape, 0.0f, 0.0f, weight_scale_, 0}, + w_down_packed); + } + scale_id_ = + AddConstInput(TensorData{TensorType_FLOAT32, scale_shape}, scale_ref_); + + out_id_ = AddOutput({GetTensorType(), out_shape}); + + std::vector empty_attrs; + flatbuffers::Offset options = + CreateStableHLOCompositeOptionsDirect( + builder_, "odml.moe_experts", + /*decomposition_subgraph_index=*/1, &empty_attrs); + + SetBuiltinOp(BuiltinOperator_STABLEHLO_COMPOSITE, + BuiltinOptions2_StableHLOCompositeOptions, options.Union()); + + BuildInterpreter({tokens_shape, rw_shape, ei_shape}, -1, false, false, + /*allocate_and_delegate=*/false); + + if (use_delegate) { + TfLiteYNNPackDelegateOptions delegate_options = + TfLiteYNNPackDelegateOptionsDefault(); + delegate_options.num_threads = 1; + delegate_options.static_shape = true; + SetDelegate(Interpreter::TfLiteDelegatePtr( + TfLiteYNNPackDelegateCreate(&delegate_options), + TfLiteYNNPackDelegateDelete)); + ApplyDelegate(); + } + if (interpreter_->AllocateTensors() != kTfLiteOk) { + fprintf(stderr, "Failed to allocate tensors\n"); + } + } + + int tokens() const { return tokens_id_; } + int rw() const { return rw_id_; } + int ei() const { return ei_id_; } + int out() const { return out_id_; } + + // The weights the delegate is expected to see, i.e. the quantized values + // scaled back to float. + const std::vector& w_gate_ref() const { return w_gate_ref_; } + const std::vector& w_up_ref() const { return w_up_ref_; } + const std::vector& w_down_ref() const { return w_down_ref_; } + const std::vector& scale_ref() const { return scale_ref_; } + + private: + // The weights as the delegate should see them: every element shares + // `weight_scale_`, so dequantizing is a single multiply. + std::vector Dequantize(const std::vector& values) const { + std::vector dequantized(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + dequantized[i] = static_cast(values[i]) * weight_scale_; + } + return dequantized; + } + + int AddScaleInput(TensorType scale_type, const std::vector& shape, + size_t count) { + if (scale_type == TensorType_FLOAT16) { + return AddConstInput(TensorData{TensorType_FLOAT16, shape}, + std::vector(count, half(weight_scale_))); + } + return AddConstInput(TensorData{TensorType_FLOAT32, shape}, + std::vector(count, weight_scale_)); + } + + const float weight_scale_; + + int tokens_id_; + int rw_id_; + int ei_id_; + int w_gate_id_; + int gate_scale_id_ = -1; + int w_up_id_; + int up_scale_id_ = -1; + int w_down_id_; + int down_scale_id_ = -1; + int scale_id_; + int out_id_; + + std::vector w_gate_ref_; + std::vector w_up_ref_; + std::vector w_down_ref_; + std::vector scale_ref_; +}; + +template +std::vector CastVector(const std::vector& values) { + std::vector cast(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + cast[i] = static_cast(values[i]); + } + return cast; +} + +template +std::vector ToFloatVector(const std::vector& values) { + std::vector as_float(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + as_float[i] = static_cast(values[i]); + } + return as_float; +} + +// Runs the quantized model and compares it against the float reference +// evaluated on the dequantized weights. The activations are quantized to int8 +// by the delegate, so the bound is relative to the magnitude of the reference +// output and has to absorb that error as well. +template +void ExpectQuantizedMoeNearReference( + int B, int N, int D_in, int D_mid, int E, int K, + const std::vector& mock_tokens, const std::vector& mock_ei, + bool explicit_scales, TensorType scale_type, TensorType weight_type, + float relative_tolerance) { + QuantizedMoeModel model(B, N, D_in, D_mid, E, K, /*use_delegate=*/true, + explicit_scales, scale_type, weight_type); + + const std::vector tokens = CastVector(mock_tokens); + const std::vector rw = + CastVector(std::vector(B * N * K, 0.5f)); + + model.template PopulateTensor(model.tokens(), tokens); + model.template PopulateTensor(model.rw(), rw); + model.PopulateTensor(model.ei(), mock_ei); + + ASSERT_EQ(model.Invoke(), kTfLiteOk); + + // The reference runs on the activations as the model stores them, so the + // bound does not have to absorb their rounding as well. + const std::vector tokens_ref = ToFloatVector(tokens); + const std::vector rw_ref = ToFloatVector(rw); + + std::vector expected_out(B * N * D_in); + RunReferenceMoe(B, N, D_in, D_mid, E, K, tokens_ref.data(), rw_ref.data(), + mock_ei.data(), model.w_gate_ref().data(), + model.w_up_ref().data(), model.w_down_ref().data(), + model.scale_ref().data(), model.scale_ref().size(), + expected_out.data()); + + float max_abs_expected = 0.0f; + for (float expected : expected_out) { + max_abs_expected = std::max(max_abs_expected, std::abs(expected)); + } + // A comparison against a reference that is ~0 would pass for any result, a + // delegate returning all zeros included. + ASSERT_GT(max_abs_expected, 1e-3f); + const float tolerance = relative_tolerance * max_abs_expected; + + const std::vector delegate_output = + ToFloatVector(model.template ExtractVector(model.out())); + for (int i = 0; i < B * N * D_in; ++i) { + EXPECT_NEAR(delegate_output[i], expected_out[i], tolerance) << "at " << i; + } +} + +// Tokens that exercise a range of magnitudes rather than a single value, so +// the dynamic quantization of the activations is not degenerate. +std::vector MakeTokens(int count) { + std::vector tokens(count); + for (int i = 0; i < count; ++i) { + tokens[i] = 0.2f * std::sin(static_cast(i)); + } + return tokens; +} + +TEST(MoeTest, QuantizedInt8MoeDecode) { + const int B = 1, N = 1, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, std::vector(B * N * D_in, 0.5f), + /*mock_ei=*/{1, 3}, /*explicit_scales=*/true, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT8, + /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, QuantizedInt8MoePrefill) { + const int B = 1, N = 4, D_in = 32, D_mid = 16, E = 4, K = 2; + std::vector mock_ei(B * N * K); + for (size_t i = 0; i < mock_ei.size(); ++i) { + mock_ei[i] = (i * 2 + 1) % E; + } + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, MakeTokens(B * N * D_in), mock_ei, + /*explicit_scales=*/true, /*scale_type=*/TensorType_FLOAT32, + /*weight_type=*/TensorType_INT8, /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, AffineQuantizedInt8MoeDecode) { + const int B = 1, N = 1, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, std::vector(B * N * D_in, 0.5f), + /*mock_ei=*/{1, 3}, /*explicit_scales=*/false, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT8, + /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, QuantizedInt4MoeDecode) { + const int B = 1, N = 1, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, std::vector(B * N * D_in, 0.5f), + /*mock_ei=*/{1, 3}, /*explicit_scales=*/true, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT4, + /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, QuantizedInt4MoePrefill) { + const int B = 1, N = 4, D_in = 32, D_mid = 16, E = 4, K = 2; + std::vector mock_ei(B * N * K); + for (size_t i = 0; i < mock_ei.size(); ++i) { + mock_ei[i] = (i * 2 + 1) % E; + } + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, MakeTokens(B * N * D_in), mock_ei, + /*explicit_scales=*/true, /*scale_type=*/TensorType_FLOAT32, + /*weight_type=*/TensorType_INT4, /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, QuantizedInt2MoeDecode) { + const int B = 1, N = 1, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, std::vector(B * N * D_in, 0.5f), + /*mock_ei=*/{1, 3}, /*explicit_scales=*/true, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT2, + /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, QuantizedInt2MoePrefill) { + const int B = 1, N = 4, D_in = 32, D_mid = 16, E = 4, K = 2; + std::vector mock_ei(B * N * K); + for (size_t i = 0; i < mock_ei.size(); ++i) { + mock_ei[i] = (i * 2 + 1) % E; + } + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, MakeTokens(B * N * D_in), mock_ei, + /*explicit_scales=*/true, /*scale_type=*/TensorType_FLOAT32, + /*weight_type=*/TensorType_INT2, /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, AffineQuantizedInt2MoeDecode) { + const int B = 1, N = 1, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, std::vector(B * N * D_in, 0.5f), + /*mock_ei=*/{1, 3}, /*explicit_scales=*/false, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT2, + /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, Float16WeightsMoe) { + int B = 1, N = 2, D_in = 16, D_mid = 32, E = 4, K = 2; + FloatMoeModel model(B, N, D_in, D_mid, E, K, + /*use_delegate=*/true); + + std::vector tokens(B * N * D_in); + std::vector rw(B * N * K); + std::vector ei = {0, 2, 1, 3}; + + for (size_t i = 0; i < tokens.size(); ++i) { + tokens[i] = 0.3f * std::sin(static_cast(i)); + } + for (size_t i = 0; i < rw.size(); ++i) { + rw[i] = 0.5f; + } + + model.PopulateTensor(model.tokens(), tokens); + model.PopulateTensor(model.rw(), rw); + model.PopulateTensor(model.ei(), ei); + + ASSERT_EQ(model.Invoke(), kTfLiteOk); + + std::vector expected_out(B * N * D_in); + RunReferenceMoe(B, N, D_in, D_mid, E, K, tokens.data(), rw.data(), ei.data(), + model.w_gate_ref().data(), model.w_up_ref().data(), + model.w_down_ref().data(), model.scale_ref().data(), + model.scale_ref().size(), expected_out.data()); + + std::vector del_out = model.ExtractVector(model.out()); + for (size_t i = 0; i < expected_out.size(); ++i) { + EXPECT_NEAR(del_out[i], expected_out[i], 1e-3f) << "Mismatch at " << i; + } +} + +TEST(MoeTest, PureFloat16Moe) { + int B = 1, N = 2, D_in = 16, D_mid = 32, E = 4, K = 2; + FloatMoeModel model(B, N, D_in, D_mid, E, K, + /*use_delegate=*/true); + + std::vector tokens(B * N * D_in); + std::vector tokens_ref(tokens.size()); + std::vector rw(B * N * K); + std::vector rw_ref(rw.size()); + std::vector ei = {0, 2, 1, 3}; + + for (size_t i = 0; i < tokens.size(); ++i) { + tokens[i] = static_cast(0.3f * std::sin(static_cast(i))); + tokens_ref[i] = static_cast(tokens[i]); + } + for (size_t i = 0; i < rw.size(); ++i) { + rw[i] = static_cast(0.5f); + rw_ref[i] = static_cast(rw[i]); + } + + model.PopulateTensor(model.tokens(), tokens); + model.PopulateTensor(model.rw(), rw); + model.PopulateTensor(model.ei(), ei); + + ASSERT_EQ(model.Invoke(), kTfLiteOk); + + std::vector expected_out(B * N * D_in); + RunReferenceMoe(B, N, D_in, D_mid, E, K, tokens_ref.data(), rw_ref.data(), + ei.data(), model.w_gate_ref().data(), model.w_up_ref().data(), + model.w_down_ref().data(), model.scale_ref().data(), + model.scale_ref().size(), expected_out.data()); + + std::vector del_out = model.ExtractVector(model.out()); + for (size_t i = 0; i < expected_out.size(); ++i) { + EXPECT_NEAR(static_cast(del_out[i]), expected_out[i], 1e-3f) + << "Mismatch at " << i; + } +} + +TEST(MoeTest, BFloat16WeightsMoe) { + int B = 1, N = 2, D_in = 16, D_mid = 32, E = 4, K = 2; + FloatMoeModel model(B, N, D_in, D_mid, E, K, + /*use_delegate=*/true); + + std::vector tokens(B * N * D_in); + std::vector rw(B * N * K); + std::vector ei = {0, 2, 1, 3}; + + for (size_t i = 0; i < tokens.size(); ++i) { + tokens[i] = 0.3f * std::sin(static_cast(i)); + } + for (size_t i = 0; i < rw.size(); ++i) { + rw[i] = 0.5f; + } + + model.PopulateTensor(model.tokens(), tokens); + model.PopulateTensor(model.rw(), rw); + model.PopulateTensor(model.ei(), ei); + + ASSERT_EQ(model.Invoke(), kTfLiteOk); + + std::vector expected_out(B * N * D_in); + RunReferenceMoe(B, N, D_in, D_mid, E, K, tokens.data(), rw.data(), ei.data(), + model.w_gate_ref().data(), model.w_up_ref().data(), + model.w_down_ref().data(), model.scale_ref().data(), + model.scale_ref().size(), expected_out.data()); + + std::vector del_out = model.ExtractVector(model.out()); + for (size_t i = 0; i < expected_out.size(); ++i) { + EXPECT_NEAR(del_out[i], expected_out[i], 1e-2f) << "Mismatch at " << i; + } +} + +TEST(MoeTest, PureBFloat16Moe) { + int B = 1, N = 2, D_in = 16, D_mid = 32, E = 4, K = 2; + FloatMoeModel model( + B, N, D_in, D_mid, E, K, /*use_delegate=*/true); + + std::vector tokens(B * N * D_in); + std::vector tokens_ref(tokens.size()); + std::vector rw(B * N * K); + std::vector rw_ref(rw.size()); + std::vector ei = {0, 2, 1, 3}; + + for (size_t i = 0; i < tokens.size(); ++i) { + tokens[i] = + static_cast(0.3f * std::sin(static_cast(i))); + tokens_ref[i] = static_cast(tokens[i]); + } + for (size_t i = 0; i < rw.size(); ++i) { + rw[i] = static_cast(0.5f); + rw_ref[i] = static_cast(rw[i]); + } + + model.PopulateTensor(model.tokens(), tokens); + model.PopulateTensor(model.rw(), rw); + model.PopulateTensor(model.ei(), ei); + + ASSERT_EQ(model.Invoke(), kTfLiteOk); + + std::vector expected_out(B * N * D_in); + RunReferenceMoe(B, N, D_in, D_mid, E, K, tokens_ref.data(), rw_ref.data(), + ei.data(), model.w_gate_ref().data(), model.w_up_ref().data(), + model.w_down_ref().data(), model.scale_ref().data(), + model.scale_ref().size(), expected_out.data()); + + std::vector del_out = + model.ExtractVector(model.out()); + for (size_t i = 0; i < expected_out.size(); ++i) { + EXPECT_NEAR(static_cast(del_out[i]), expected_out[i], 1e-2f) + << "Mismatch at " << i; + } +} + +TEST(MoeTest, QuantizedInt8MoeFloat16Scale) { + const int B = 1, N = 2, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, MakeTokens(B * N * D_in), + /*mock_ei=*/{0, 2, 1, 3}, /*explicit_scales=*/true, + /*scale_type=*/TensorType_FLOAT16, /*weight_type=*/TensorType_INT8, + /*relative_tolerance=*/0.01f); +} + +// 16-bit activations are quantized to int8 directly, with nothing promoting +// them to fp32 first. +TEST(MoeTest, QuantizedInt8MoeFloat16Activations) { + const int B = 1, N = 2, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, MakeTokens(B * N * D_in), + /*mock_ei=*/{0, 2, 1, 3}, /*explicit_scales=*/true, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT8, + /*relative_tolerance=*/0.01f); +} + +TEST(MoeTest, QuantizedInt8MoeBFloat16Activations) { + const int B = 1, N = 2, D_in = 32, D_mid = 16, E = 4, K = 2; + ExpectQuantizedMoeNearReference( + B, N, D_in, D_mid, E, K, MakeTokens(B * N * D_in), + /*mock_ei=*/{0, 2, 1, 3}, /*explicit_scales=*/true, + /*scale_type=*/TensorType_FLOAT32, /*weight_type=*/TensorType_INT8, + /*relative_tolerance=*/0.02f); +} + } // namespace } // namespace ynnpack } // namespace tflite diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py index 5298c63077fca5..9665e7917c50bf 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -2104,6 +2104,19 @@ def testUpperBoundFloatHandCoded(self): tf_result = self.evaluate(array_ops.searchsorted(cdf, arr, side="right")) self.assertAllEqual(result, tf_result) + def testScalarValue(self): + cdf = np.array([0, 0.2, 0.5, 0.6, 0.8, 1.0], dtype=np.float32) + values = (0.53, np.float32(0.53), constant_op.constant(0.53)) + + for value in values: + for side in ("left", "right"): + with self.subTest(value=value, side=side): + result = np.searchsorted(cdf, 0.53, side=side) + tf_result = self.evaluate( + array_ops.searchsorted(cdf, value, side=side) + ) + self.assertAllEqual(result, tf_result) + def testUpperBoundFloatRandomNd(self): dim_size = 7 for d in range(1, 5): diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 13c62bb1a55042..225111e25efc70 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -6131,6 +6131,12 @@ def searchsorted(sorted_sequence, If the first `N-1` dimensions of the two tensors don't match. """ sequence_size = shape_internal(sorted_sequence)[-1] + values = ops.convert_to_tensor(values, name="values") + values_shape = shape_internal(values) + + if values.shape.rank == 0: + values = expand_dims(values, 0) + values_size = shape_internal(values)[-1] sorted_sequence_2d = reshape(sorted_sequence, [-1, sequence_size]) values_2d = reshape(values, [-1, values_size]) @@ -6143,7 +6149,7 @@ def searchsorted(sorted_sequence, else: raise ValueError("Argument `side` must be either 'right' or 'left'. " f"Received: `side` = '{side}'.") - return reshape(output, shape_internal(values)) + return reshape(output, values_shape) quantize.__doc__ = gen_array_ops.quantize_v2.__doc__ diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index bfb5ccbf5c70df..9811798a45bd13 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -1326,6 +1326,20 @@ def stack(arrays, axis=0): # pylint: disable=missing-function-docstring unwrapped_arrays = [ a if isinstance(a, np_arrays.ndarray) else a for a in arrays ] + # NumPy raises AxisError for out-of-bounds axes instead of letting the + # backend kernel fail with a confusing error. `axis` is an insertion + # position, so rank itself is in bounds (NumPy 2.x allows axis == rank). + if arrays: + maybe_rank = asarray(unwrapped_arrays[0]).shape.rank + if ( + maybe_rank is not None + and isinstance(axis, (int, np.integer)) + and not -maybe_rank - 1 <= axis <= maybe_rank + ): + raise ValueError( + f'Argument `axis` (received axis={axis}) is out of bounds ' + f'for input of rank {maybe_rank}.' + ) return asarray(array_ops_stack.stack(unwrapped_arrays, axis)) diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py index 054549a78fea00..a971d815f73551 100644 --- a/tensorflow/python/ops/numpy_ops/tests/np_test.py +++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py @@ -39,7 +39,14 @@ config.parse_flags_with_absl() -nonempty_nonscalar_array_shapes = [(4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)] +nonempty_nonscalar_array_shapes = [ + (4,), + (3, 4), + (3, 1), + (1, 4), + (2, 1, 4), + (2, 3, 4), +] nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes empty_array_shapes = [(0,), (0, 4), (3, 0),] @@ -2958,6 +2965,19 @@ def foo(x): return tnp.concatenate(x) foo(onp.zeros((2, 2))) # doesn't crash + def testStackInvalidAxis(self): + # NumPy raises AxisError for out-of-bounds stack axes. `axis` is an + # insertion position, so rank itself is in bounds. + a = onp.ones((2, 3)) + with self.assertRaisesRegex(ValueError, "out of bounds"): + tnp.stack([a, a], axis=3) + with self.assertRaisesRegex(ValueError, "out of bounds"): + tnp.stack([a, a], axis=-4) + with self.assertRaisesRegex(ValueError, "out of bounds"): + tnp.stack([a, a], axis=4) + # In-bounds boundaries still work. + self.assertEqual(tnp.stack([a, a], axis=2).shape, (2, 3, 2)) + @jtu.disable def testReluGradientConstants(self): # This is a regression test that verifies that constants associated with the diff --git a/tensorflow/python/ops/ragged/ragged_cross_op_test.py b/tensorflow/python/ops/ragged/ragged_cross_op_test.py index ce3dc913f35e3d..e37c017df3c73a 100644 --- a/tensorflow/python/ops/ragged/ragged_cross_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_cross_op_test.py @@ -446,6 +446,29 @@ def testSparseValuesAndIndicesMustMatch(self): out_values_type=dtypes.string, out_row_splits_type=dtypes.int64)) + def testSparseIndicesAndValuesRowCountMustMatch(self): + # A sparse input whose values are longer than its indices matrix used to + # walk the indices tensor out of bounds while building the row splits. + with self.assertRaisesRegex( + (ValueError, errors.InvalidArgumentError), 'Expected size of values' + ): + self.evaluate( + gen_ragged_array_ops.RaggedCross( + ragged_values=[], + ragged_row_splits=[], + sparse_indices=[[[0, 0]]], + sparse_values=[['a', 'b', 'c']], + sparse_shape=[[1, 10]], + dense_inputs=[], + input_order='S', + hashed_output=False, + num_buckets=0, + hash_key=2, + out_values_type=dtypes.string, + out_row_splits_type=dtypes.int64, + ) + ) + def testRaggedValuesAndSplitsMustMatch(self): with self.assertRaisesRegex( (ValueError, errors.InvalidArgumentError), diff --git a/tensorflow/tf_version.bzl b/tensorflow/tf_version.bzl index 10f296197af5ab..64e53b9d7c0c32 100644 --- a/tensorflow/tf_version.bzl +++ b/tensorflow/tf_version.bzl @@ -23,7 +23,7 @@ load( # These constants are used by the targets //third_party/tensorflow/core/public:release_version, # //third_party/tensorflow:tensorflow_bzl and //third_party/tensorflow/tools/pip_package:setup_py. -TF_VERSION = "2.22.0" +TF_VERSION = "2.23.0" MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION = TF_VERSION.split(".") TF_WHEEL_VERSION_SUFFIX = VERSION_SUFFIX TF_SEMANTIC_VERSION_SUFFIX = SEMANTIC_VERSION_SUFFIX diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 882b2410fbb926..6a7be06dcfc0ce 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -46,9 +46,9 @@ bazel_dep(name = "rules_ml_toolchain") # echo "sha256-${HASH}" archive_override( module_name = "rules_ml_toolchain", - integrity = "sha256-MC3tNh8VTCDFB7VH7tMT8B1Ozlerxr5m8Pd7EMEHZCk=", - strip_prefix = "rules_ml_toolchain-2d318f0d6a4c7a402d30a0bce260040e6807b196", - urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/2d318f0d6a4c7a402d30a0bce260040e6807b196.tar.gz"], + integrity = "sha256-XrL/AKDPawv6pfBMKS1QlceKHPIu1OXi0wfH42qmfe8=", + strip_prefix = "rules_ml_toolchain-cdfa52fa7e8bd9fb51cf66de1da166969d850edd", + urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/cdfa52fa7e8bd9fb51cf66de1da166969d850edd.tar.gz"], ) # TODO: Upstream the patch? @@ -303,6 +303,13 @@ use_repo(local_clang_configure, "local_config_clang") rocm_configure = use_extension("//third_party/extensions:rocm_configure.bzl", "rocm_configure_ext") use_repo(rocm_configure, "config_rocm_hipcc", "local_config_rocm") +rocm_hermetic = use_extension("@rules_ml_toolchain//extensions:rocm_hermetic_download.bzl", "rocm_hermetic_download_ext") + +override_repo( + rocm_hermetic, + config_rocm_hipcc = "config_rocm_hipcc", +) + tensorrt_configure = use_extension("//third_party/extensions:tensorrt_configure.bzl", "tensorrt_configure_ext") use_repo(tensorrt_configure, "local_config_tensorrt") diff --git a/third_party/xla/build_tools/lint/run_dwyu_test.py b/third_party/xla/build_tools/lint/run_dwyu_test.py index ec3d8a39cd84da..2cf877d02f4fc4 100644 --- a/third_party/xla/build_tools/lint/run_dwyu_test.py +++ b/third_party/xla/build_tools/lint/run_dwyu_test.py @@ -378,6 +378,114 @@ def test_absl_strings_and_string_view_filtering(self): ) self.assertEqual(result.stdout, expected_stdout) + def _setup_googletest_repo(self): + self.write_file( + "MODULE.bazel", + """\ +module(name = "xla") +bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest") +""", + ) + gtest_repo = self.external / "googletest+" + (gtest_repo / "googletest/include/gtest").mkdir(parents=True) + (gtest_repo / "googlemock/include/gmock").mkdir(parents=True) + (gtest_repo / "googlemock/src").mkdir(parents=True) + (gtest_repo / "MODULE.bazel").write_text('module(name = "googletest")\n') + (gtest_repo / "googletest/include/gtest/gtest.h").write_text("// gtest\n") + (gtest_repo / "googlemock/include/gmock/gmock.h").write_text("// gmock\n") + (gtest_repo / "googlemock/src/gmock_main.cc").write_text("int main() {}\n") + (gtest_repo / "BUILD.bazel").write_text("""\ +cc_library( + name = "gtest", + hdrs = [ + "googlemock/include/gmock/gmock.h", + "googletest/include/gtest/gtest.h", + ], + includes = [ + "googlemock", + "googlemock/include", + "googletest", + "googletest/include", + ], + visibility = ["//visibility:public"], +) + +cc_library( + name = "gtest_main", + srcs = ["googlemock/src/gmock_main.cc"], + hdrs = [ + "googlemock/include/gmock/gmock.h", + "googletest/include/gtest/gtest.h", + ], + includes = [ + "googlemock", + "googlemock/include", + "googletest", + "googletest/include", + ], + tags = [ + "avoid_dep", + "keep_dep", + ], + visibility = ["//visibility:public"], + deps = [":gtest"], +) +""") + + def test_googletest_gtest_main_compatibility(self): + self._setup_googletest_repo() + test_cases = [ + ( + "gtest_main_satisfies_quoted_gtest_and_gmock_includes", + '["@com_google_googletest//:gtest_main"]', + '#include "gmock/gmock.h"\n#include "gtest/gtest.h"\n', + 0, + "", + ), + ( + "gtest_main_kept_for_angle_bracket_and_transitive_includes", + '["//xla/platform:errors", "@com_google_googletest//:gtest_main"]', + ( + "#include \n" + "#include \n" + '#include "xla/platform/errors.h"\n' + ), + 0, + "", + ), + ( + "missing_gtest_suggests_gtest_not_gtest_main", + "[]", + '#include "gtest/gtest.h"\n', + 3, + ( + "buildozer 'add deps @com_google_googletest//:gtest'" + " @xla//xla/consumer:consumer_test\n" + ), + ), + ] + for name, deps, source, expected_returncode, expected_stdout in test_cases: + with self.subTest(name): + self.write_file( + "xla/consumer/BUILD", + f"""\ +cc_test( + name = "consumer_test", + srcs = ["consumer_test.cc"], + deps = {deps}, +) +""", + ) + self.write_file("xla/consumer/consumer_test.cc", source) + result = self.check_targets("//xla/consumer:consumer_test") + self.assertEqual( + result.returncode, + expected_returncode, + result.stdout + result.stderr, + ) + self.assertEqual(result.stdout, expected_stdout) + if __name__ == "__main__": unittest.main() + diff --git a/third_party/xla/third_party/googletest/googletest.patch b/third_party/xla/third_party/googletest/googletest.patch index 6d7063653eba8d..740bfef363fd1e 100644 --- a/third_party/xla/third_party/googletest/googletest.patch +++ b/third_party/xla/third_party/googletest/googletest.patch @@ -16,7 +16,7 @@ diff --git a/BUILD.bazel b/BUILD.bazel index cc254457..49120384 100644 --- a/BUILD.bazel +++ b/BUILD.bazel -@@ -178,6 +178,10 @@ alias( +@@ -181,6 +181,24 @@ alias( cc_library( name = "gtest_main", srcs = ["googlemock/src/gmock_main.cc"], @@ -24,6 +24,20 @@ index cc254457..49120384 100644 + "googletest/include/gtest/*.h", + "googlemock/include/gmock/*.h", + ]), ++ includes = [ ++ "googlemock", ++ "googlemock/include", ++ "googletest", ++ "googletest/include", ++ ], ++ tags = [ ++ # Prefer `:gtest` when DWYU suggests adding a missing gtest dependency. ++ "avoid_dep", ++ # Prevent DWYU from removing `gtest_main` (needed for its `main()` entry point). ++ # Note: Do not set `alwayslink = True` here, as targets that also link ++ # `//tensorflow/core:test_main` would get duplicate `main()` definitions (ODR violation). ++ "keep_dep", ++ ], features = select({ ":windows": ["windows_export_all_symbols"], "//conditions:default": [], diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl index 3f9a1f215739ac..87c10ae95de02d 100644 --- a/third_party/xla/workspace3.bzl +++ b/third_party/xla/workspace3.bzl @@ -74,10 +74,10 @@ def workspace(): # Details: https://github.com/google-ml-infra/rules_ml_toolchain tf_http_archive( name = "rules_ml_toolchain", - sha256 = "302ded361f154c20c507b547eed313f01d4ece57abc6be66f0f77b10c1076429", - strip_prefix = "rules_ml_toolchain-2d318f0d6a4c7a402d30a0bce260040e6807b196", + sha256 = "5eb2ff00a0cf6b0bfaa5f04c292d5095c78a1cf22ed4e5e2d307c7e36aa67def", + strip_prefix = "rules_ml_toolchain-cdfa52fa7e8bd9fb51cf66de1da166969d850edd", urls = tf_mirror_urls( - "https://github.com/google-ml-infra/rules_ml_toolchain/archive/2d318f0d6a4c7a402d30a0bce260040e6807b196.tar.gz", + "https://github.com/google-ml-infra/rules_ml_toolchain/archive/cdfa52fa7e8bd9fb51cf66de1da166969d850edd.tar.gz", ), ) diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo_benchmark_runner.cc b/third_party/xla/xla/backends/cpu/benchmarks/hlo_benchmark_runner.cc index b0950394aa2dd6..8d57db79d494dd 100644 --- a/third_party/xla/xla/backends/cpu/benchmarks/hlo_benchmark_runner.cc +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo_benchmark_runner.cc @@ -174,6 +174,9 @@ absl::Status RunHloBenchmarkImpl(benchmark::State* absl_nullable state, compile_options.executable_build_options.mutable_debug_options() ->add_xla_disable_hlo_passes("cpu-parallel-task-assigner"); } + // Disable HLO module upload for benchmarks. + compile_options.executable_build_options.mutable_debug_options() + ->set_xla_enable_hlo_modules_upload(false); std::unique_ptr executable; if (benchmark_options.aot_options) { diff --git a/third_party/xla/xla/backends/cpu/codegen/cpu_features.cc b/third_party/xla/xla/backends/cpu/codegen/cpu_features.cc index a259bd1d0846e2..d21b324b5671cc 100644 --- a/third_party/xla/xla/backends/cpu/codegen/cpu_features.cc +++ b/third_party/xla/xla/backends/cpu/codegen/cpu_features.cc @@ -57,6 +57,8 @@ absl::string_view CpuTargetFromMaxFeature(CPUFeature max_feature) { return "sapphirerapids"; case CPUFeature::AMX_FP16: return "graniterapids"; + case CPUFeature::AMX_FP8: + return "diamondrapids"; //===------------------------------------------------------------------===// // AArch64 @@ -90,6 +92,7 @@ std::optional CpuFeatureFromString(absl::string_view cpu_feature) { {"AVX512_BF16", CPUFeature::AVX512_BF16}, {"AMX", CPUFeature::AMX_BF16}, // Includes AMX_INT8. {"AMX_FP16", CPUFeature::AMX_FP16}, + {"AMX_FP8", CPUFeature::AMX_FP8}, //===-------------------------------------------------------------===// // AArch64 //===-------------------------------------------------------------===// @@ -158,6 +161,19 @@ static bool ShouldEnableX86CpuFeature(absl::string_view feature, if (feature == "amx-fp16") return false; [[fallthrough]]; + case CPUFeature::AMX_FP16: + // Suppress DMR-and-newer additions over GNR: AVX10.2, remaining AMX + // tiles, and APX extensions. + if (feature == "amx-fp8" || feature == "amx-tf32" || + feature == "amx-avx512" || feature == "amx-movrs" || + absl::StartsWith(feature, "avx10.2") || feature == "egpr" || + feature == "ndd" || feature == "ccmp" || feature == "nf" || + feature == "zu" || feature == "ppx" || feature == "push2pop2" || + feature == "movrs" || feature == "jmpabs") { + return false; + } + [[fallthrough]]; + default: // Leave all other features enabled. return true; diff --git a/third_party/xla/xla/backends/cpu/codegen/target_machine_features.cc b/third_party/xla/xla/backends/cpu/codegen/target_machine_features.cc index 4c426a18c12cd6..244d9b6e78a917 100644 --- a/third_party/xla/xla/backends/cpu/codegen/target_machine_features.cc +++ b/third_party/xla/xla/backends/cpu/codegen/target_machine_features.cc @@ -37,10 +37,12 @@ TargetMachineFeatures::TargetMachineFeatures( // Calling virtual methods in the constructor is discouraged, so we don't // call `get_target_feature_string` here. if (target_machine_) { - has_avx512bf16_ = absl::StrContains( - target_machine_->getTargetFeatureString().str(), "+avx512bf16"); - has_avx512fp16_ = absl::StrContains( - target_machine_->getTargetFeatureString().str(), "+avx512fp16"); + const std::string feature_string = + target_machine_->getTargetFeatureString().str(); + has_avx512bf16_ = absl::StrContains(feature_string, "+avx512bf16"); + has_avx512fp16_ = absl::StrContains(feature_string, "+avx512fp16"); + has_amx_fp16_ = absl::StrContains(feature_string, "+amx-fp16"); + has_amx_fp8_ = absl::StrContains(feature_string, "+amx-fp8"); } } diff --git a/third_party/xla/xla/backends/cpu/codegen/target_machine_features.h b/third_party/xla/xla/backends/cpu/codegen/target_machine_features.h index a0c338447919b9..0f044e50442b09 100644 --- a/third_party/xla/xla/backends/cpu/codegen/target_machine_features.h +++ b/third_party/xla/xla/backends/cpu/codegen/target_machine_features.h @@ -73,6 +73,8 @@ class TargetMachineFeatures { virtual bool has_avx512bf16() const { return has_avx512bf16_; } virtual bool has_avx512fp16() const { return has_avx512fp16_; } + virtual bool has_amx_fp16() const { return has_amx_fp16_; } + virtual bool has_amx_fp8() const { return has_amx_fp8_; } private: llvm::TargetTransformInfo* GetTargetTransformInfoFor( @@ -87,6 +89,8 @@ class TargetMachineFeatures { // Store availability of popular features here for efficient checks. bool has_avx512bf16_ = false; bool has_avx512fp16_ = false; + bool has_amx_fp16_ = false; + bool has_amx_fp8_ = false; }; } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/target_machine_features_test.cc b/third_party/xla/xla/backends/cpu/codegen/target_machine_features_test.cc index 54757c444b77d4..023711bdf664db 100644 --- a/third_party/xla/xla/backends/cpu/codegen/target_machine_features_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/target_machine_features_test.cc @@ -65,5 +65,51 @@ INSTANTIATE_TEST_SUITE_P(Avx512Bf16Fp16Suite, Avx512Bf16Fp16Test, ::testing::ValuesIn(GetAvx512Bf16Fp16TestSpecs()), Avx512Bf16Fp16Test::Name); +struct AmxFp16Fp8TestSpec { + std::string cpu_name; + std::string features; + bool has_amx_fp16; + bool has_amx_fp8; +}; + +class AmxFp16Fp8Test + : public TargetMachineTestBase, + public ::testing::WithParamInterface { + public: + static std::string Name( + const ::testing::TestParamInfo& info) { + return info.param.cpu_name; + } +}; + +TEST_P(AmxFp16Fp8Test, CheckAvailability) { + AmxFp16Fp8TestSpec spec = GetParam(); + const char* triple_string = "x86_64-unknown-linux-gnu"; + std::unique_ptr features = + CreateTargetMachineFeatures(triple_string, spec.cpu_name, spec.features); + EXPECT_EQ(features->has_amx_fp16(), spec.has_amx_fp16); + EXPECT_EQ(features->has_amx_fp8(), spec.has_amx_fp8); +} + +std::vector GetAmxFp16Fp8TestSpecs() { + return std::vector{ + // SPR-class: neither AMX_FP16 nor AMX_FP8. + AmxFp16Fp8TestSpec{"sapphirerapids", "+amx-bf16,+amx-int8,+amx-tile", + false, false}, + // GNR-class: AMX_FP16 only -> emulation path for FP8. + AmxFp16Fp8TestSpec{"graniterapids", + "+amx-bf16,+amx-int8,+amx-tile,+amx-fp16", true, + false}, + // DMR-class: native AMX_FP8 (and AMX_FP16). + AmxFp16Fp8TestSpec{"diamondrapids", + "+amx-bf16,+amx-int8,+amx-tile,+amx-fp16,+amx-fp8", + true, true}, + }; +} + +INSTANTIATE_TEST_SUITE_P(AmxFp16Fp8Suite, AmxFp16Fp8Test, + ::testing::ValuesIn(GetAmxFp16Fp8TestSpecs()), + AmxFp16Fp8Test::Name); + } // namespace } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/onednn_support.cc b/third_party/xla/xla/backends/cpu/onednn_support.cc index b1ec97c5947cd8..cc712b01d03b72 100644 --- a/third_party/xla/xla/backends/cpu/onednn_support.cc +++ b/third_party/xla/xla/backends/cpu/onednn_support.cc @@ -46,9 +46,29 @@ bool IsOneDnnSupportedDType(PrimitiveType dtype, // Check for data type support if target machine features are provided. // Unit tests may provide target machine features to simulate different CPU // capabilities. - return (cpu_features != nullptr && - ((dtype == BF16 && cpu_features->has_avx512bf16()) || - (dtype == F16 && cpu_features->has_avx512fp16()))); + if (cpu_features == nullptr) { + return false; + } + switch (dtype) { + case BF16: + return cpu_features->has_avx512bf16(); + case F16: + return cpu_features->has_avx512fp16(); + case F8E5M2: + case F8E4M3FN: + case F8E4M3: + if (cpu_features->has_amx_fp8()) { + return true; + } + if (cpu_features->has_amx_fp16()) { + LOG_FIRST_N(INFO, 1) << "XLA:CPU FP8 dispatched via oneDNN AMX-FP16 " + "emulation path (target ISA lacks AMX-FP8)."; + return true; + } + return false; + default: + return false; + } } bool IsOneDnnSupportedLayout(const Shape& shape) { diff --git a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc index cf0b7ffcb587f6..79afc3055657ea 100644 --- a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc +++ b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher_test.cc @@ -31,6 +31,7 @@ class YnnE2eTest : public HloTestBase { debug_options.add_xla_cpu_experimental_ynn_fusion_type( DebugOptions::LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION); debug_options.clear_xla_cpu_experimental_ynn_fusion_type(); + debug_options.set_xla_cpu_experimental_onednn_custom_call(false); return debug_options; } }; @@ -60,6 +61,7 @@ class YnnReduceTest : public HloTestBase { DebugOptions debug_options = HloTestBase::GetDebugOptionsForTest(); debug_options.add_xla_cpu_experimental_ynn_fusion_type( DebugOptions::LIBRARY_FUSION_TYPE_REDUCE); + debug_options.set_xla_cpu_experimental_onednn_custom_call(false); return debug_options; } }; @@ -218,6 +220,7 @@ class YnnDotTest : public HloTestBase, DebugOptions debug_options = HloTestBase::GetDebugOptionsForTest(); debug_options.add_xla_cpu_experimental_ynn_fusion_type( DebugOptions::LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT); + debug_options.set_xla_cpu_experimental_onednn_custom_call(false); return debug_options; } }; @@ -279,6 +282,7 @@ class YnnReduceEltwiseTest : public HloTestBase { DebugOptions::LIBRARY_FUSION_TYPE_REDUCE); debug_options.add_xla_cpu_experimental_ynn_fusion_type( DebugOptions::LIBRARY_FUSION_TYPE_ELTWISE); + debug_options.set_xla_cpu_experimental_onednn_custom_call(false); return debug_options; } }; diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD index 83186a4ab519b3..fec89081d96b7b 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/BUILD +++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD @@ -68,7 +68,6 @@ cc_library( tags = ["gpu"], deps = [ ":gpu_codegen_backend", - "//xla:autotuning_proto_cc", "//xla:xla_proto_cc", "//xla/backends/autotuner:backends_proto_cc", "//xla/backends/autotuner:codegen_backend", @@ -87,8 +86,6 @@ cc_library( "//xla/service/gpu/model:gpu_indexing_performance_model", "//xla/stream_executor:device_description", "//xla/stream_executor/gpu:tma_metadata", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:inlined_vector", @@ -121,7 +118,6 @@ xla_test( use_legacy_runtime = True, deps = [ ":block_level_emitter", - "//xla:autotuning_proto_cc", "//xla:debug_options_flags", "//xla:xla_proto_cc", "//xla/backends/autotuner:codegen_backend", @@ -134,11 +130,11 @@ xla_test( "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:ir_emission_utils", "//xla/service/gpu:nvptx_compiler_impl", - "//xla/stream_executor:device_description_proto_cc", + "//xla/stream_executor:platform", "//xla/stream_executor:stream_executor_h", - "//xla/stream_executor/gpu:tma_metadata", "//xla/tsl/platform:statusor", "//xla/tsl/util/proto:proto_matchers", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest_main", @@ -712,6 +708,7 @@ cc_library( "//xla/hlo/analysis:alias_info", "//xla/service:compiler", "//xla/service:hlo_cost_analysis", + "//xla/stream_executor:device_address_allocator", "//xla/stream_executor:stream_executor_h", "@com_google_absl//absl/types:span", "@llvm-project//mlir:IR", @@ -1048,7 +1045,6 @@ xla_test( "//xla/hlo/ir:hlo", "//xla/service:compiler", "//xla/service:platform_util", - "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_compiler", "//xla/service/gpu/autotuning:config_assigner_pass", "//xla/stream_executor:device_address_allocator", @@ -1136,6 +1132,8 @@ xla_test( "//xla/stream_executor:stream_executor_memory_allocator", "//xla/stream_executor/platform:platform_object_registry", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest_main", "@llvm-project//mlir:IR", diff --git a/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc b/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc index 4a8dc055887e27..787ebad44635c4 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc @@ -54,7 +54,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/service/compiler.h" #include "xla/service/gpu/autotuning/config_assigner_pass.h" -#include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_compiler.h" #include "xla/service/platform_util.h" #include "xla/status_macros.h" diff --git a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc index deb9c45372d554..8329f81c5bceb4 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.cc @@ -30,7 +30,6 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" -#include "xla/autotuning.pb.h" #include "xla/backends/autotuner/codegen_backend.h" #include "xla/backends/gpu/codegen/triton/support.h" #include "xla/backends/gpu/codegen/triton/tma_utils.h" @@ -46,8 +45,6 @@ limitations under the License. #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" #include "xla/stream_executor/gpu/tma_metadata.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "xla/xla.pb.h" #include "triton/Version.h" diff --git a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc index 60f09e6e1215c7..4dc530e3e877e5 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc @@ -21,9 +21,9 @@ limitations under the License. #include #include +#include "absl/log/check.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" -#include "xla/autotuning.pb.h" #include "xla/backends/autotuner/codegen_backend.h" #include "xla/codegen/xtile/xtile_config.pb.h" #include "xla/debug_options_flags.h" @@ -35,8 +35,7 @@ limitations under the License. #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/nvptx_compiler.h" #include "xla/service/platform_util.h" -#include "xla/stream_executor/device_description.pb.h" -#include "xla/stream_executor/gpu/tma_metadata.h" +#include "xla/stream_executor/platform.h" #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/platform/statusor.h" #include "xla/tsl/util/proto/proto_matchers.h" @@ -58,6 +57,16 @@ bool AnyTmaAllowed(const std::vector>& configs) { }); } +// Returns the first device executor of the default platform. +se::StreamExecutor* GetDefaultStreamExecutor() { + absl::StatusOr platform = PlatformUtil::GetDefaultPlatform(); + CHECK_OK(platform.status()); + absl::StatusOr stream_executor = + (*platform)->ExecutorForDevice(0); + CHECK_OK(stream_executor.status()); + return *stream_executor; +} + // Test fixture for the TritonBlockLevelFusionEmitterBackend. // // Inherits from HloHardwareIndependentTestBase to use XLA utilities like @@ -68,10 +77,7 @@ class TritonBlockLevelFusionEmitterBackendTest protected: TritonBlockLevelFusionEmitterBackendTest() : debug_options_(GetDebugOptionsFromFlags()), - stream_executor_(PlatformUtil::GetDefaultPlatform() - .value() - ->ExecutorForDevice(0) - .value()), + stream_executor_(GetDefaultStreamExecutor()), target_config_(stream_executor_), backend_(&debug_options_, &compiler_, compiler_.ShapeSizeBytesFunction(), &target_config_) {} diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory.h b/third_party/xla/xla/backends/gpu/autotuner/factory.h index bf7cc43a365748..29ec8fce04ee78 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/factory.h +++ b/third_party/xla/xla/backends/gpu/autotuner/factory.h @@ -27,6 +27,7 @@ limitations under the License. #include "xla/hlo/analysis/alias_info.h" #include "xla/service/compiler.h" #include "xla/service/hlo_cost_analysis.h" +#include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/stream_executor.h" namespace xla { diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc index 6883cb33d534d0..b9cd546490dcd5 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc @@ -16,9 +16,13 @@ limitations under the License. #include "xla/backends/gpu/autotuner/factory.h" #include +#include +#include #include #include +#include "absl/log/check.h" +#include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "mlir/IR/MLIRContext.h" #include "xla/backends/autotuner/backends.pb.h" @@ -50,6 +54,32 @@ struct FactoryTestParams { bool run_on_rocm = true; }; +// Returns the canonical GPU platform. +se::Platform* GetGpuPlatform() { + absl::StatusOr name = PlatformUtil::CanonicalPlatformName("gpu"); + CHECK_OK(name.status()); + absl::StatusOr platform = + se::PlatformManager::PlatformWithName(absl::AsciiStrToUpper(*name)); + CHECK_OK(platform.status()); + return *platform; +} + +// Returns the compiler registered for `platform`. +std::unique_ptr GetCompilerForPlatform(se::Platform* platform) { + absl::StatusOr> compiler = + Compiler::GetForPlatform(platform->id()); + CHECK_OK(compiler.status()); + return *std::move(compiler); +} + +// Returns the first device executor of `platform`. +se::StreamExecutor* GetStreamExecutor(se::Platform* platform) { + absl::StatusOr stream_executor = + platform->ExecutorForDevice(0); + CHECK_OK(stream_executor.status()); + return *stream_executor; +} + class FactoryTest : public xla::HloHardwareIndependentTestBase, public ::testing::WithParamInterface { protected: @@ -61,12 +91,9 @@ class FactoryTest : public xla::HloHardwareIndependentTestBase, stream_executor::StreamExecutorAddressAllocator allocator_; FactoryTest() - : platform_(se::PlatformManager::PlatformWithName( - absl::AsciiStrToUpper( - PlatformUtil::CanonicalPlatformName("gpu").value())) - .value()), - compiler_(xla::Compiler::GetForPlatform(platform_->id()).value()), - stream_executor_(platform_->ExecutorForDevice(0).value()), + : platform_(GetGpuPlatform()), + compiler_(GetCompilerForPlatform(platform_)), + stream_executor_(GetStreamExecutor(platform_)), target_config_(stream_executor_), allocator_(stream_executor_) {} }; diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc index 4c12bc5b0afbf1..2f763269e6c645 100644 --- a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc @@ -137,9 +137,8 @@ class CuDnnFusionFileCheckTest : public CuDnnFusionTest { const std::string root_name( module->entry_computation()->root_instruction()->name()); BinaryMap dnn_compiled_graphs; - CuDnnFusionCompiler cudnn_compiler(stream_executor()->AsDnn(), - se::DeviceDescription(), - dnn_compiled_graphs); + CuDnnFusionCompiler cudnn_compiler( + stream_executor()->AsDnn(), device_description(), dnn_compiled_graphs); // Run filecheck even if CuDnnFusionCompiler failed. cudnn_compiler.Run(module.get()).IgnoreError(); std::string dump; @@ -283,8 +282,8 @@ ENTRY e { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(kHloText)); BinaryMap dnn_compiled_graphs; - CuDnnFusionCompiler cudnn_compiler( - stream_executor()->AsDnn(), se::DeviceDescription(), dnn_compiled_graphs); + CuDnnFusionCompiler cudnn_compiler(stream_executor()->AsDnn(), + device_description(), dnn_compiled_graphs); ASSERT_OK_AND_ASSIGN(bool changed, cudnn_compiler.Run(module.get())); EXPECT_TRUE(changed); EXPECT_THAT(module->entry_computation()->root_instruction(), @@ -317,8 +316,8 @@ e { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(kHloText)); BinaryMap dnn_compiled_graphs; - CuDnnFusionCompiler cudnn_compiler( - stream_executor()->AsDnn(), se::DeviceDescription(), dnn_compiled_graphs); + CuDnnFusionCompiler cudnn_compiler(stream_executor()->AsDnn(), + device_description(), dnn_compiled_graphs); EXPECT_THAT(cudnn_compiler.Run(module.get()), absl_testing::IsOkAndHolds(false)); // Single dot is not supported by cuDNN, so Triton should be used. @@ -366,8 +365,8 @@ ENTRY e { ROOT r = tuple(f0, f1) })")); BinaryMap dnn_compiled_graphs; - CuDnnFusionCompiler cudnn_compiler( - stream_executor()->AsDnn(), se::DeviceDescription(), dnn_compiled_graphs); + CuDnnFusionCompiler cudnn_compiler(stream_executor()->AsDnn(), + device_description(), dnn_compiled_graphs); ASSERT_OK_AND_ASSIGN(bool changed, cudnn_compiler.Run(module.get())); EXPECT_TRUE(changed); EXPECT_THAT(module->entry_computation()->root_instruction(), @@ -594,7 +593,7 @@ ENTRY e { ASSERT_OK_AND_ASSIGN(std::unique_ptr module_live, ParseAndReturnVerifiedModule(kHlo)); BinaryMap binary_map_live; - CuDnnFusionCompiler live_compiler(executor->AsDnn(), se::DeviceDescription(), + CuDnnFusionCompiler live_compiler(executor->AsDnn(), device_description, binary_map_live); ASSERT_OK_AND_ASSIGN(bool changed_live, live_compiler.Run(module_live.get())); ASSERT_TRUE(changed_live); diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 6a8bb809caa9dd..fcb31cc5e44121 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -3258,7 +3258,6 @@ xla_cc_test( srcs = ["softmax_rewriter_triton_test.cc"], deps = [ ":softmax_rewriter_triton", - "//xla:xla_data_proto_cc", "//xla:xla_proto_cc", "//xla/backends/gpu/codegen/triton:support", "//xla/hlo/analysis:symbolic_map", diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2.cc index 8e5e982a830343..b6a43b9c91ea8c 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2.cc @@ -143,6 +143,23 @@ bool HasSupportedShapes(const HloInstruction* hero) { return false; } } + absl::flat_hash_set seen_tuple_indices; + for (const HloInstruction* user : hero->users()) { + const auto* gte = DynCast(user); + if (gte == nullptr) { + LOG(WARNING) << "DynamicSliceFusionRewriterV2: skipping " + << hero->name() + << " because tuple result has non-GTE user " + << user->name(); + return false; + } + if (!seen_tuple_indices.insert(gte->tuple_index()).second) { + LOG(WARNING) << "DynamicSliceFusionRewriterV2: skipping " + << hero->name() << " because tuple result has duplicate " + << "GTE user for index " << gte->tuple_index(); + return false; + } + } } return true; } @@ -310,6 +327,9 @@ HloInstruction* FindGte(HloInstruction* hero, int64_t tuple_index) { // Walks forward from `gte` to find a DUS chain. If found, returns a // SlicedResult with the GTE prepended into noops. std::optional ResolveLeafDus(HloInstruction* gte) { + if (gte->user_count() != 1) { + return std::nullopt; + } for (HloInstruction* user : gte->users()) { if (auto sliced = ResolveSlicedResult(user)) { sliced->noops.insert(sliced->noops.begin(), gte); @@ -322,6 +342,9 @@ std::optional ResolveLeafDus(HloInstruction* gte) { // Resolves the sliced result for a non-tuple hero. std::vector ResolveNonTupleSlicedResult( HloInstruction* hero, const CaptureUpdateSlice& capture_update_slice) { + if (hero->user_count() != 1) { + return {}; + } for (HloInstruction* user : hero->users()) { if (auto sliced = ResolveSlicedResult(user)) { if (!capture_update_slice(hero, std::nullopt, sliced->update_slice)) { @@ -689,6 +712,10 @@ absl::StatusOr RewriteHero( plan->external_operands, fusion_body)); module->SetAndUniquifyInstrName(fusion, "dynamic_slice_fusion"); ABSL_RETURN_IF_ERROR(SetDynamicSliceFusionBackendConfig(fusion)); + ABSL_RETURN_IF_ERROR(fusion->CopyAllControlDepsFrom(hero)); + ABSL_RETURN_IF_ERROR(hero->DropAllControlDeps()); + + const bool hero_has_side_effect = hero->HasSideEffect(); if (sliced_results.size() > 1) { bool any_result_replaced = false; @@ -707,11 +734,22 @@ absl::StatusOr RewriteHero( } if (!any_result_replaced) { ABSL_RETURN_IF_ERROR(parent->ReplaceInstruction(hero, fusion)); + } else if (hero_has_side_effect) { + ABSL_RETURN_IF_ERROR(parent->RemoveInstructionAndUnusedOperands(hero)); } } else if (sliced_results.size() == 1) { if (sliced_results[0].update_slice != nullptr) { ABSL_RETURN_IF_ERROR( parent->ReplaceInstruction(sliced_results[0].update_slice, fusion)); + if (hero_has_side_effect) { + ABSL_RETURN_IF_ERROR(parent->RemoveInstructionAndUnusedOperands(hero)); + } + } else if (hero->shape().IsTuple() && !sliced_results[0].noops.empty()) { + ABSL_RETURN_IF_ERROR( + parent->ReplaceInstruction(sliced_results[0].noops.back(), fusion)); + if (hero_has_side_effect) { + ABSL_RETURN_IF_ERROR(parent->RemoveInstructionAndUnusedOperands(hero)); + } } else { ABSL_RETURN_IF_ERROR(parent->ReplaceInstruction(hero, fusion)); } diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2_test.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2_test.cc index 4349e45e8effd8..75137c28b06b54 100644 --- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2_test.cc @@ -2546,5 +2546,266 @@ TEST_F(DynamicSliceFusionRewriterV2Test, O2LooksThroughOptBarrier) { fusion_checks); } +TEST_F(DynamicSliceFusionRewriterV2Test, SideEffectingTupleHeroRemoved) { + const char* hlo = R"( + HloModule test + + body { + p0 = (s32[], f32[4,8,8], f32[8,8], f32[8,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + input = f32[4,8,8] get-tuple-element(p0), index=1 + prev0 = f32[8,8] get-tuple-element(p0), index=2 + prev1 = f32[8,8] get-tuple-element(p0), index=3 + c0 = s32[] constant(0) + ds = f32[1,8,8] dynamic-slice(input, ivar, c0, c0), + dynamic_slice_sizes={1,8,8} + bc = f32[8,8] bitcast(ds) + hero = (f32[8,8], f32[8,8]) custom-call(bc), + custom_call_target="fake_target", + custom_call_has_side_effect=true + gte0 = f32[8,8] get-tuple-element(hero), index=0 + gte1 = f32[8,8] get-tuple-element(hero), index=1 + c1 = s32[] constant(1) + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], f32[4,8,8], f32[8,8], f32[8,8]) + tuple(next_ivar, input, gte0, gte1) + } + + condition { + p0 = (s32[], f32[4,8,8], f32[8,8], f32[8,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c4 = s32[] constant(4) + ROOT cmp = pred[] compare(ivar, c4), direction=LT + } + + ENTRY main { + input = f32[4,8,8] parameter(0) + prev0 = f32[8,8] parameter(1) + prev1 = f32[8,8] parameter(2) + c0 = s32[] constant(0) + tuple = (s32[], f32[4,8,8], f32[8,8], f32[8,8]) + tuple(c0, input, prev0, prev1) + ROOT while = (s32[], f32[4,8,8], f32[8,8], f32[8,8]) + while(tuple), condition=condition, body=body, + backend_config={"known_trip_count":{"n":"4"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + } + )"; + + const char* expected = R"( + ; CHECK: %dynamic-slice-fusion{{.*}} { + ; CHECK: {{.*}} dynamic-slice( + ; CHECK: {{.*}} bitcast( + ; CHECK: {{.*}} custom-call( + ; CHECK-SAME: custom_call_target="fake_target" + ; CHECK-SAME: custom_call_has_side_effect=true + ; CHECK: {{.*}} get-tuple-element( + ; CHECK: {{.*}} get-tuple-element( + ; CHECK: ROOT {{.*}} tuple( + ; CHECK: } + ; CHECK: body + ; CHECK-NOT: dynamic-slice( + ; CHECK-NOT: custom-call( + ; CHECK: {{.*}} fusion( + ; CHECK-SAME: kind=kCustom + ; CHECK-SAME: "name":"dynamic_slice_fusion" + ; CHECK-NOT: custom-call( + ; CHECK: } + )"; + + RunAndFilecheckHloRewrite(hlo, MakePipeline(), expected); +} + +TEST_F(DynamicSliceFusionRewriterV2Test, + SideEffectingSingleOutputDusHeroRemoved) { + const char* hlo = R"( + HloModule test + + body { + p0 = (s32[], f32[4,8,8], f32[8,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + buf = f32[4,8,8] get-tuple-element(p0), index=1 + in = f32[8,8] get-tuple-element(p0), index=2 + c0 = s32[] constant(0) + hero = f32[8,8] custom-call(in), + custom_call_target="fake_target", + custom_call_has_side_effect=true + bc = f32[1,8,8] bitcast(hero) + dus = f32[4,8,8] dynamic-update-slice(buf, bc, ivar, c0, c0) + c1 = s32[] constant(1) + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], f32[4,8,8], f32[8,8]) tuple(next_ivar, dus, in) + } + + condition { + p0 = (s32[], f32[4,8,8], f32[8,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c4 = s32[] constant(4) + ROOT cmp = pred[] compare(ivar, c4), direction=LT + } + + ENTRY main { + buf = f32[4,8,8] parameter(0) + in = f32[8,8] parameter(1) + c0 = s32[] constant(0) + tuple = (s32[], f32[4,8,8], f32[8,8]) tuple(c0, buf, in) + ROOT while = (s32[], f32[4,8,8], f32[8,8]) while(tuple), + condition=condition, body=body, + backend_config={"known_trip_count":{"n":"4"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + } + )"; + + const char* expected = R"( + ; CHECK: %dynamic-slice-fusion{{.*}} { + ; CHECK: {{.*}} custom-call( + ; CHECK-SAME: custom_call_target="fake_target" + ; CHECK-SAME: custom_call_has_side_effect=true + ; CHECK: {{.*}} bitcast( + ; CHECK: ROOT {{.*}} dynamic-update-slice( + ; CHECK: } + ; CHECK: body + ; CHECK-NOT: custom-call( + ; CHECK: {{.*}} fusion( + ; CHECK-SAME: kind=kCustom + ; CHECK-SAME: "name":"dynamic_slice_fusion" + ; CHECK-NOT: custom-call( + ; CHECK: } + )"; + + RunAndFilecheckHloRewrite(hlo, MakePipeline(), expected); +} + +TEST_F(DynamicSliceFusionRewriterV2Test, + HeroWithDusAndExternalUserNotDuplicated) { + // When a non-tuple hero with no sliced inputs feeds both a DUS and an + // external user, it must not absorb the DUS and leave a duplicate hero. + const char* hlo = R"( + HloModule test + + body { + param = (s32[], f32[64], f32[4,64]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + res = f32[64] get-tuple-element(param), index=1 + buf = f32[4,64] get-tuple-element(param), index=2 + hero = f32[64] custom-call(res), custom_call_target="fake_target" + hero_2d = f32[1,64] bitcast(hero) + zero = s32[] constant(0) + updated = f32[4,64] dynamic-update-slice(buf, hero_2d, i, zero) + new_res = f32[64] add(res, hero) + one = s32[] constant(1) + next_i = s32[] add(i, one) + ROOT tuple = (s32[], f32[64], f32[4,64]) tuple(next_i, new_res, updated) + } + + cond { + param = (s32[], f32[64], f32[4,64]) parameter(0) + i = s32[] get-tuple-element(param), index=0 + limit = s32[] constant(4) + ROOT cmp = pred[] compare(i, limit), direction=LT + } + + ENTRY main { + res = f32[64] parameter(0) + buf = f32[4,64] parameter(1) + zero = s32[] constant(0) + init = (s32[], f32[64], f32[4,64]) tuple(zero, res, buf) + ROOT while = (s32[], f32[64], f32[4,64]) while(init), + condition=cond, body=body, + backend_config={"known_trip_count":{"n":"4"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + } + )"; + + const char* expected = R"( + ; CHECK-NOT: dynamic-slice-fusion + ; CHECK: body + ; CHECK: {{.*}} custom-call( + ; CHECK-NOT: custom-call( + )"; + + RunAndFilecheckHloRewrite(hlo, MakePipeline(), expected); +} + +TEST_F(DynamicSliceFusionRewriterV2Test, + TupleOutputDusLeafWithExternalUserBecomesPassthrough) { + // Hero produces (f32[8,8], f32[8,8]) from a sliced input. Output 0 feeds + // both a DUS and an external user (`add`), while output 1 feeds only a DUS. + // Output 0 must become a passthrough output (replacing all uses of `gte0`) + // and output 1's DUS is fused, leaving no duplicate hero in `body`. + const char* hlo = R"( + HloModule test + + body { + p0 = (s32[], f32[4,8,8], f32[4,8,8], f32[4,8,8], f32[8,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + in = f32[4,8,8] get-tuple-element(p0), index=1 + buf0 = f32[4,8,8] get-tuple-element(p0), index=2 + buf1 = f32[4,8,8] get-tuple-element(p0), index=3 + accum = f32[8,8] get-tuple-element(p0), index=4 + c0 = s32[] constant(0) + ds = f32[1,8,8] dynamic-slice(in, ivar, c0, c0), + dynamic_slice_sizes={1,8,8} + bc_in = f32[8,8] bitcast(ds) + hero = (f32[8,8], f32[8,8]) custom-call(bc_in), + custom_call_target="fake_target" + gte0 = f32[8,8] get-tuple-element(hero), index=0 + gte1 = f32[8,8] get-tuple-element(hero), index=1 + bc0 = f32[1,8,8] bitcast(gte0) + dus0 = f32[4,8,8] dynamic-update-slice(buf0, bc0, ivar, c0, c0) + new_accum = f32[8,8] add(accum, gte0) + bc1 = f32[1,8,8] bitcast(gte1) + dus1 = f32[4,8,8] dynamic-update-slice(buf1, bc1, ivar, c0, c0) + c1 = s32[] constant(1) + next_ivar = s32[] add(ivar, c1) + ROOT result = (s32[], f32[4,8,8], f32[4,8,8], f32[4,8,8], f32[8,8]) + tuple(next_ivar, in, dus0, dus1, new_accum) + } + + condition { + p0 = (s32[], f32[4,8,8], f32[4,8,8], f32[4,8,8], f32[8,8]) parameter(0) + ivar = s32[] get-tuple-element(p0), index=0 + c4 = s32[] constant(4) + ROOT cmp = pred[] compare(ivar, c4), direction=LT + } + + ENTRY main { + in = f32[4,8,8] parameter(0) + buf0 = f32[4,8,8] parameter(1) + buf1 = f32[4,8,8] parameter(2) + accum = f32[8,8] parameter(3) + c0 = s32[] constant(0) + tuple = (s32[], f32[4,8,8], f32[4,8,8], f32[4,8,8], f32[8,8]) + tuple(c0, in, buf0, buf1, accum) + ROOT while = (s32[], f32[4,8,8], f32[4,8,8], f32[4,8,8], f32[8,8]) + while(tuple), condition=condition, body=body, + backend_config={"known_trip_count":{"n":"4"}, + "known_init_step":{"init":"0","step":"1"}, + "known_induction_variable":{"tuple_index":"0"}} + } + )"; + + const char* expected = R"( + ; CHECK: %dynamic-slice-fusion{{.*}} { + ; CHECK: {{.*}} dynamic-slice( + ; CHECK: {{.*}} custom-call( + ; CHECK: {{.*}} dynamic-update-slice( + ; CHECK: ROOT {{.*}} = (f32[8,8]{1,0}, f32[4,8,8]{2,1,0}) tuple( + ; CHECK: } + ; CHECK: body + ; CHECK-NOT: custom-call( + ; CHECK: {{.*}} fusion( + ; CHECK-SAME: kind=kCustom + ; CHECK-SAME: "name":"dynamic_slice_fusion" + ; CHECK-NOT: custom-call( + ; CHECK: } + )"; + + RunAndFilecheckHloRewrite(hlo, MakePipeline(), expected); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.cc b/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.cc index 487599bdef510e..4e3b10ffa18f94 100644 --- a/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.cc +++ b/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.cc @@ -95,11 +95,10 @@ static absl::StatusOr AsynchronizeInstruction(HloInstruction* instr) { } ClearSchedulingAnnotations(instr); - ABSL_ASSIGN_OR_RETURN( - HloInstruction * done, - computation->CreateAsyncInstructions( - instr, {}, ExplicitStreamAnnotationAsyncWrapper::kMainExecutionThread, - /*replace=*/true)); + ABSL_ASSIGN_OR_RETURN(HloInstruction * done, + computation->CreateAsyncInstructions( + instr, {}, HloInstruction::kParallelExecutionThread, + /*replace=*/true)); // Replace the original attributes after creating the async pair. done->set_frontend_attributes(original_attributes); done->mutable_operand(0)->set_frontend_attributes(original_attributes); diff --git a/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.h b/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.h index 57d0bc8716475c..0e509ac487d753 100644 --- a/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.h +++ b/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper.h @@ -32,8 +32,6 @@ namespace xla::gpu { // asynchronized, using AsyncStart/Done pairs to achieve asynchronous execution. class ExplicitStreamAnnotationAsyncWrapper : public HloModulePass { public: - inline static constexpr char kMainExecutionThread[] = "main"; - absl::string_view name() const override { return "explicit-stream-annotation-async-wrapper"; } diff --git a/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper_test.cc b/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper_test.cc index 4273f442f47993..34c982d2e80367 100644 --- a/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/explicit_stream_annotation_async_wrapper_test.cc @@ -21,8 +21,10 @@ limitations under the License. #include #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" @@ -60,7 +62,7 @@ TEST_F(ExplicitStreamAnnotationAsyncWrapperTest, AnnotatedOpIsWrapped) { TF_ASSERT_OK_AND_ASSIGN(bool mutated, wrapper_pass.Run(module.get())); absl::StatusOr filecheck_result = RunFileCheck(module->ToString({}), R"( // CHECK: %lhs.1 = f32[] constant(42) - // CHECK: %call-start = ((f32[]), f32[]) call-start(%lhs.1), to_apply=%sub, frontend_attributes={_xla_stream_annotation="1"} + // CHECK: %call-start = ((f32[]), f32[]) call-start(%lhs.1), async_execution_thread="parallel", to_apply=%sub, frontend_attributes={_xla_stream_annotation="1"} // CHECK: ROOT %call-done = f32[] call-done(%call-start), frontend_attributes={_xla_stream_annotation="1"}, backend_config={"operation_queue_id":"0","force_earliest_schedule":false )"); ASSERT_OK(filecheck_result.status()); @@ -105,9 +107,9 @@ TEST_F(ExplicitStreamAnnotationAsyncWrapperTest, OverlappingGemms) { ASSERT_TRUE(mutated); absl::StatusOr filecheck_result = RunFileCheck(module->ToString({}), R"( - // CHECK: %call-start = ((f32[2048,2048]{1,0}, f32[2048,2048]{1,0}), f32[2048,2048]{1,0}) call-start(%x, %y), to_apply=%gemm1, frontend_attributes={_scheduling_group_id="0",_xla_stream_annotation="2"} + // CHECK: %call-start = ((f32[2048,2048]{1,0}, f32[2048,2048]{1,0}), f32[2048,2048]{1,0}) call-start(%x, %y), async_execution_thread="parallel", to_apply=%gemm1, frontend_attributes={_scheduling_group_id="0",_xla_stream_annotation="2"} // CHECK: %call-done = f32[2048,2048]{1,0} call-done(%call-start), frontend_attributes={_scheduling_group_id="0",_xla_stream_annotation="2"}, backend_config={"operation_queue_id":"0","force_earliest_schedule":false - // CHECK: %call-start.1 = ((f32[2048,2048]{1,0}, f32[2048,2048]{1,0}), f32[2048,2048]{1,0}) call-start(%x, %y), to_apply=%gemm2, frontend_attributes={_scheduling_group_id="1",_xla_stream_annotation="1"} + // CHECK: %call-start.1 = ((f32[2048,2048]{1,0}, f32[2048,2048]{1,0}), f32[2048,2048]{1,0}) call-start(%x, %y), async_execution_thread="parallel", to_apply=%gemm2, frontend_attributes={_scheduling_group_id="1",_xla_stream_annotation="1"} // CHECK: ROOT %call-done.1 = f32[2048,2048]{1,0} call-done(%call-start.1), frontend_attributes={_scheduling_group_id="1",_xla_stream_annotation="1"}, backend_config={"operation_queue_id":"0","force_earliest_schedule":false )"); ASSERT_OK(filecheck_result.status()); @@ -173,7 +175,7 @@ TEST_F(ExplicitStreamAnnotationAsyncWrapperTest, ASSERT_TRUE(mutated); absl::StatusOr filecheck_result = RunFileCheck(module->ToString({}), R"( - // CHECK: %call-start = {{.*}} call-start(%a), to_apply= + // CHECK: %call-start = {{.*}} call-start(%a), async_execution_thread="parallel", to_apply= // CHECK-SAME: frontend_attributes={_xla_stream_annotation="1"} // CHECK: ROOT %call-done = f32[4]{0} call-done(%call-start) // CHECK-SAME: frontend_attributes={_xla_stream_annotation="1"} @@ -434,5 +436,51 @@ TEST_F(ExplicitStreamAnnotationAsyncWrapperTest, EXPECT_NE(FindInstruction(module.get(), "cc"), nullptr); } +TEST_F(ExplicitStreamAnnotationAsyncWrapperTest, + AsyncPairUsesParallelExecutionThread) { + // Verify that the async-start/done pair created by the wrapper uses + // async_execution_thread="parallel" (StreamAttributeAsyncWrapper:: + // kParallelExecutionThread) rather than the old "main" thread. The parallel + // thread is what causes ExecutionStreamAssignment to route the op to a + // communication stream so it can overlap with compute ops on the main stream. + const absl::string_view hlo_string = R"( + HloModule m + + ENTRY %main (a: f32[4]) -> f32[4] { + %a = f32[4]{0} parameter(0) + ROOT %cc = f32[4]{0} custom-call(%a), custom_call_target="te_op", + frontend_attributes={_xla_stream_annotation="collective",inlineable="false"} + })"; + + auto debug_options = HloHardwareIndependentTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_stream_annotation(true); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + module->mutable_config().set_debug_options(debug_options); + ExplicitStreamAnnotationAsyncWrapper wrapper_pass; + + ASSERT_OK_AND_ASSIGN(bool mutated, wrapper_pass.Run(module.get())); + ASSERT_TRUE(mutated); + + // The printed form of async-start only emits async_execution_thread when it + // differs from "main", so presence of the attribute confirms the change. + absl::StatusOr filecheck_result = RunFileCheck(module->ToString({}), R"( + // CHECK: call-start{{.*}}async_execution_thread="parallel" + )"); + ASSERT_OK(filecheck_result.status()); + EXPECT_TRUE(*filecheck_result); + + // Also verify programmatically: the start instruction's execution thread and + // the wrapped computation's thread must both equal kParallelExecutionThread. + const HloInstruction* call_start = + FindInstruction(module.get(), "call-start"); + ASSERT_NE(call_start, nullptr); + const auto* async_start = DynCast(call_start); + ASSERT_NE(async_start, nullptr); + EXPECT_EQ(async_start->async_execution_thread(), + HloInstruction::kParallelExecutionThread); + EXPECT_EQ(async_start->async_wrapped_computation()->execution_thread(), + HloInstruction::kParallelExecutionThread); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc index 518e99ab073e3b..6af936e36b83f7 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc @@ -42,7 +42,6 @@ limitations under the License. #include "xla/stream_executor/device_description.h" #include "xla/tsl/platform/errors.h" #include "xla/xla.pb.h" -#include "xla/xla_data.pb.h" namespace xla { namespace gpu { diff --git a/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.cc b/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.cc index 98c11519557e60..15a6af535de935 100644 --- a/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.cc +++ b/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.cc @@ -44,11 +44,10 @@ static absl::StatusOr AsynchronizeInstruction(HloInstruction* instr) { return false; } HloComputation* computation = instr->parent(); - ABSL_ASSIGN_OR_RETURN( - HloInstruction * done, - computation->CreateAsyncInstructions( - instr, {}, StreamAttributeAsyncWrapper::kParallelExecutionThread, - /*replace=*/true)); + ABSL_ASSIGN_OR_RETURN(HloInstruction * done, + computation->CreateAsyncInstructions( + instr, {}, HloInstruction::kParallelExecutionThread, + /*replace=*/true)); ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config, done->backend_config()); // Set the false delay of done op to be false so it can be scheduled diff --git a/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.h b/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.h index 107f40cef1f684..9929fbb4fc602a 100644 --- a/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.h +++ b/third_party/xla/xla/backends/gpu/transforms/stream_attribute_async_wrapper.h @@ -32,8 +32,6 @@ namespace xla::gpu { // asynchronous executions. class StreamAttributeAsyncWrapper : public HloModulePass { public: - inline static constexpr char kParallelExecutionThread[] = "parallel"; - absl::string_view name() const override { return "async-stream-attribute-wrapper"; } diff --git a/third_party/xla/xla/codegen/emitters/fusion_wrapper_base.cc b/third_party/xla/xla/codegen/emitters/fusion_wrapper_base.cc index efa3c357fe03d0..bc108af646486b 100644 --- a/third_party/xla/xla/codegen/emitters/fusion_wrapper_base.cc +++ b/third_party/xla/xla/codegen/emitters/fusion_wrapper_base.cc @@ -56,6 +56,8 @@ absl::StatusOr FusionWrapperBase::RunImpl( computation->AddInstruction(HloInstruction::CreateFusion( instruction->shape(), ChooseFusionKind(*instruction, *instruction), instruction)); + fusion_instruction->fused_instructions_computation()->SetExecutionThread( + computation->execution_thread()); const absl::string_view wrapped_opcode = HloOpcodeString(instruction->opcode()); module->SetAndUniquifyInstrName(fusion_instruction, diff --git a/third_party/xla/xla/codegen/xtile/codegen/BUILD b/third_party/xla/xla/codegen/xtile/codegen/BUILD index c2873d4854359c..b3264481aefb8b 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/BUILD +++ b/third_party/xla/xla/codegen/xtile/codegen/BUILD @@ -246,6 +246,7 @@ xla_cc_test( srcs = ["tiled_emitter_constraints_test.cc"], deps = [ ":tiled_emitter_constraints", + "//xla:xla_proto_cc", "//xla/codegen/tiling:symbolic_tile_analysis", "//xla/codegen/tiling:tiling_specification", "//xla/hlo/analysis:symbolic_map", diff --git a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc index c49fdef1078e33..863b83b4af1ec1 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc +++ b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/service/instruction_fusion.h" #include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" +#include "xla/xla.pb.h" namespace xla { namespace { diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index d2db95a8c0da98..f9f29c90840802 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -411,7 +411,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_triton_gemm_any(true); opts.set_xla_gpu_experimental_gemm_fusion_v2(false); opts.set_xla_gpu_verify_triton_fusion_numerics(false); - opts.set_xla_gpu_experimental_enable_tiling_propagation(false); + opts.set_xla_gpu_experimental_enable_tiling_propagation(true); opts.set_xla_gpu_experimental_cost_model_gemm_tiling_default(false); // Moving reduce-scatter out of while loops can increase memory footprint, so @@ -1756,8 +1756,8 @@ void MakeDebugOptionsFlags(std::vector* flag_list, debug_options->xla_cpu_max_isa(), "Maximum ISA that XLA:CPU LLVM backend will codegen, i.e., it will not " "use newer instructions. Available values: SSE4_2, AVX, AVX2, AVX512, " - "AVX512_VNNI, AVX512_BF16, AMX, and AMX_FP16. (`AMX` will enable both " - "`AMX_BF16` and `AMX_INT8` instructions.)")); + "AVX512_VNNI, AVX512_BF16, AMX, AMX_FP16, and AMX_FP8. (`AMX` will " + "enable both `AMX_BF16` and `AMX_INT8` instructions.)")); flag_list->push_back(tsl::Flag( "xla_cpu_emitter_verification_level", int32_setter_for(&DebugOptions::set_xla_cpu_emitter_verification_level), diff --git a/third_party/xla/xla/debug_options_parsers_test.cc b/third_party/xla/xla/debug_options_parsers_test.cc index 4c8fc76c0ffee4..78d35a6bd31642 100644 --- a/third_party/xla/xla/debug_options_parsers_test.cc +++ b/third_party/xla/xla/debug_options_parsers_test.cc @@ -110,6 +110,7 @@ std::vector GetUppercaseStringSetterTestCases() { UppercaseStringSetterTestSpec{"sse4_2", "SSE4_2"}, UppercaseStringSetterTestSpec{"aVx512", "AVX512"}, UppercaseStringSetterTestSpec{"AMx_fP16", "AMX_FP16"}, + UppercaseStringSetterTestSpec{"amx_FP8", "AMX_FP8"}, }); } diff --git a/third_party/xla/xla/ffi/api/record_api.h b/third_party/xla/xla/ffi/api/record_api.h index 16acc890115a8e..554b9409b9d5ec 100644 --- a/third_party/xla/xla/ffi/api/record_api.h +++ b/third_party/xla/xla/ffi/api/record_api.h @@ -21,7 +21,6 @@ limitations under the License. #include #include #include -#include #include #include @@ -51,6 +50,7 @@ using KernelArg = std::variant; enum class SourceFormat { kPtx = XLA_FFI_SourceFormat_PTX, kCubin = XLA_FFI_SourceFormat_CUBIN, + kFunctionPtr = XLA_FFI_SourceFormat_FUNCTION_PTR, }; enum class RecordAction { @@ -122,7 +122,7 @@ class RecordContextBase { // Converts a span of `KernelArg` or `void*` to a vector of // `XLA_FFI_KernelArg`. template - std::vector ConvertArgs(KernelArgSpan args) { + std::vector ConvertArgs(const KernelArgSpan& args) { std::vector raw_args; const size_t num_args = std::size(args); raw_args.reserve(num_args); @@ -175,8 +175,13 @@ class RecordContextBase { StatusOr CreateLaunch( const char* kernel_name, const void* kernel_data, size_t kernel_size, SourceFormat format, XLA_FFI_LaunchDims launch_dims, - uint32_t shared_mem_bytes, bool uses_pdl, KernelArgSpan args, - DepSpan dependencies = {}) { + uint32_t shared_mem_bytes, bool uses_pdl, const KernelArgSpan& args, + const DepSpan& dependencies = {}) { + if (kernel_data == nullptr) { + return ErrorPolicy::FromErrorCode( + XLA_FFI_Error_Code_INVALID_ARGUMENT, + "kernel_data/function_ptr must not be null during create_launch."); + } std::vector raw_args = ConvertArgs(args); XLA_FFI_KernelArgs ffi_args{raw_args.data(), static_cast(raw_args.size())}; @@ -207,7 +212,8 @@ class RecordContextBase { } template - Status UpdateLaunch(const XLA_FFI_Command* command, KernelArgSpan args) { + Status UpdateLaunch(const XLA_FFI_Command* command, + const KernelArgSpan& args) { std::vector raw_args = ConvertArgs(args); XLA_FFI_KernelArgs ffi_args{raw_args.data(), static_cast(raw_args.size())}; diff --git a/third_party/xla/xla/ffi/api/record_c_api.h b/third_party/xla/xla/ffi/api/record_c_api.h index 260950ceb09ea6..2ad49f23fa5183 100644 --- a/third_party/xla/xla/ffi/api/record_c_api.h +++ b/third_party/xla/xla/ffi/api/record_c_api.h @@ -41,7 +41,7 @@ extern "C" { //============================================================================== #define XLA_FFI_Extension_Record 128 #define XLA_FFI_Extension_Record_MajorVersion 0 -#define XLA_FFI_Extension_Record_MinorVersion 1 +#define XLA_FFI_Extension_Record_MinorVersion 2 //===----------------------------------------------------------------------===// // Command Buffer Recording API (FFI Record) @@ -64,6 +64,7 @@ typedef enum XLA_FFI_RecordAction { typedef enum XLA_FFI_SourceFormat { XLA_FFI_SourceFormat_PTX = 0, XLA_FFI_SourceFormat_CUBIN = 1, + XLA_FFI_SourceFormat_FUNCTION_PTR = 2, } XLA_FFI_SourceFormat; typedef struct XLA_FFI_Dim3 { @@ -142,9 +143,11 @@ typedef struct XLA_FFI_KernelArgs { // ``` typedef struct XLA_FFI_RecordApi { // Creates a launch command for a kernel with the given name, data and size. - // kernel_data is the binary of the kernel, and format is the format of the - // kernel data. Since recording implies creation of a graph, dependencies - // argument is used to specify dependencies on other commands. + // kernel_data is the binary of the kernel (for PTX/CUBIN) or the pre-loaded + // device function pointer (for FUNCTION_PTR, where kernel_size must be 0), + // and format is the format of the kernel data. Since recording implies + // creation of a graph, dependencies argument is used to specify dependencies + // on other commands. // For eg: // XLA_FFI_Command* cmd1 = api->create_launch(...); // XLA_FFI_Command* cmd2 = api->create_launch(..., {&cmd1}, 1); diff --git a/third_party/xla/xla/ffi/api/record_ffi_test.cc b/third_party/xla/xla/ffi/api/record_ffi_test.cc index 6ee2c4569d754c..f0250de167809a 100644 --- a/third_party/xla/xla/ffi/api/record_ffi_test.cc +++ b/third_party/xla/xla/ffi/api/record_ffi_test.cc @@ -42,8 +42,10 @@ TEST(FfiTest, RecordStage) { EXPECT_EQ(ctx.commands().capacity(), 10); EXPECT_EQ(ctx.commands().size(), 0); auto cmd = ctx.CreateLaunch( - "test", nullptr, 0, SourceFormat::kPtx, {{1, 1, 1}, {1, 1, 1}}, 0, - /*uses_pdl=*/false, std::vector{}); + /*kernel_name=*/"test", /*kernel_data=*/"dummy_ptx", + /*kernel_size=*/9, /*format=*/SourceFormat::kPtx, + /*launch_dims=*/{{1, 1, 1}, {1, 1, 1}}, /*shared_mem_bytes=*/0, + /*uses_pdl=*/false, /*args=*/std::vector{}); return Error::Success(); }); @@ -65,11 +67,11 @@ TEST(FfiTest, RecordStage) { int64_t num_commands = 0; XLA_FFI_RecordFrame record_frame = {/*record_ctx=*/nullptr, - &mock_api, - XLA_FFI_RecordAction_Create, - commands_storage, - &num_commands, - 10}; + /*api=*/&mock_api, + /*action=*/XLA_FFI_RecordAction_Create, + /*commands=*/commands_storage, + /*num_commands=*/&num_commands, + /*max_commands=*/10}; RecordExtension::CExtension record_frame_ext = BuildRecordCExtension(&record_frame); @@ -84,5 +86,111 @@ TEST(FfiTest, RecordStage) { EXPECT_TRUE(called); } +TEST(FfiTest, CreateLaunchWithFunctionPtr) { + struct CaptureState { + const void* kernel_data = nullptr; + int64_t kernel_size = -1; + XLA_FFI_SourceFormat format = XLA_FFI_SourceFormat_PTX; + const XLA_FFI_Command* dummy_cmd = + reinterpret_cast(0x1); + } state; + + const void* expected_func_ptr = reinterpret_cast(0xDEADBEEF); + + auto handler = Ffi::BindRecord().Ctx>().To( + [&](RecordContext ctx) { + auto null_cmd = ctx.CreateLaunch( + /*kernel_name=*/"null_func", /*kernel_data=*/nullptr, + /*kernel_size=*/0, /*format=*/SourceFormat::kFunctionPtr, + /*launch_dims=*/{{1, 1, 1}, {1, 1, 1}}, /*shared_mem_bytes=*/0, + /*uses_pdl=*/false, /*args=*/std::vector{}); + EXPECT_FALSE(null_cmd.has_value()); + + auto cmd = ctx.CreateLaunch( + /*kernel_name=*/"test_func_ptr", /*kernel_data=*/expected_func_ptr, + /*kernel_size=*/0, /*format=*/SourceFormat::kFunctionPtr, + /*launch_dims=*/{{1, 1, 1}, {1, 1, 1}}, /*shared_mem_bytes=*/0, + /*uses_pdl=*/false, /*args=*/std::vector{}); + EXPECT_TRUE(cmd.has_value()); + EXPECT_EQ(*cmd, state.dummy_cmd); + return Error::Success(); + }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + static constexpr auto kMockApi = [](auto... args) -> XLA_FFI_Error* { + return nullptr; + }; + XLA_FFI_RecordApi mock_api{ + /*.create_launch=*/+[](XLA_FFI_RecordContext* ctx, const char*, + const void* kernel_data, int64_t kernel_size, + XLA_FFI_SourceFormat format, XLA_FFI_LaunchDims, + uint32_t, int32_t, const XLA_FFI_KernelArgs*, + const XLA_FFI_Command* const*, uint32_t, + const XLA_FFI_Command** out_command) + -> XLA_FFI_Error* { + auto* s = reinterpret_cast(ctx); + s->kernel_data = kernel_data; + s->kernel_size = kernel_size; + s->format = format; + *out_command = s->dummy_cmd; + return nullptr; + }, + /*.update_launch=*/kMockApi, + /*.create_memcpy_d2d=*/kMockApi, + /*.update_memcpy_d2d=*/kMockApi, + /*.request_stream_capture=*/kMockApi, + /*.create_empty_command=*/kMockApi, + }; + + const XLA_FFI_Command* commands_storage[10] = {nullptr}; + int64_t num_commands = 0; + + XLA_FFI_RecordFrame record_frame = { + /*record_ctx=*/reinterpret_cast(&state), + /*api=*/&mock_api, + /*action=*/XLA_FFI_RecordAction_Create, + /*commands=*/commands_storage, + /*num_commands=*/&num_commands, + /*max_commands=*/10}; + + RecordExtension::CExtension record_frame_ext = + BuildRecordCExtension(&record_frame); + + InvokeContext context; + context.extension_start = &record_frame_ext.extension_base; + + ASSERT_OK( + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kRecord)); + EXPECT_EQ(state.kernel_data, expected_func_ptr); + EXPECT_EQ(state.kernel_size, 0); + EXPECT_EQ(state.format, XLA_FFI_SourceFormat_FUNCTION_PTR); + EXPECT_EQ(num_commands, 1); +} + +TEST(FfiTest, RecordExtensionVersionMismatch) { + auto handler = Ffi::BindRecord().Ctx>().To( + [&](RecordContext) { return Error::Success(); }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + XLA_FFI_RecordFrame record_frame = {}; + RecordExtension::CExtension record_frame_ext = + BuildRecordCExtension(&record_frame); + record_frame_ext.extension_base.id.minor_version = 1; + + InvokeContext context; + context.extension_start = &record_frame_ext.extension_base; + + auto status = + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kRecord); + EXPECT_FALSE(status.ok()); + EXPECT_THAT( + status.message(), + testing::HasSubstr("Extension version mismatch for RecordExtension")); +} + } // namespace } // namespace xla::ffi diff --git a/third_party/xla/xla/ffi/record_ffi_test.cc b/third_party/xla/xla/ffi/record_ffi_test.cc index d76b68f47c4fe1..99668c67f1ebef 100644 --- a/third_party/xla/xla/ffi/record_ffi_test.cc +++ b/third_party/xla/xla/ffi/record_ffi_test.cc @@ -111,13 +111,16 @@ TEST(FfiTest, RecordStageCreateAndUpdateCommands) { if (action == RecordAction::kCreate) { XLA_FFI_LaunchDims dims{{1, 1, 1}, {1, 1, 1}}; auto launch_cmd = ctx.CreateLaunch( - "dummy_kernel", nullptr, 0, SourceFormat::kPtx, dims, 0, - /*uses_pdl=*/false, std::initializer_list{}); + /*kernel_name=*/"dummy_kernel", /*kernel_data=*/"dummy_ptx", + /*kernel_size=*/9, /*format=*/SourceFormat::kPtx, + /*launch_dims=*/dims, /*shared_mem_bytes=*/0, + /*uses_pdl=*/false, /*args=*/std::initializer_list{}); EXPECT_THAT(launch_cmd, IsOkAndHolds(Eq(state.launch_cmd))); state.launch_created = true; int src = 0, dst = 0; - auto memcpy_cmd = ctx.CreateMemcpyD2D(&dst, &src, sizeof(int)); + auto memcpy_cmd = ctx.CreateMemcpyD2D( + /*dst=*/&dst, /*src=*/&src, /*size=*/sizeof(int)); EXPECT_THAT(memcpy_cmd, IsOkAndHolds(Eq(state.memcpy_cmd))); state.memcpy_created = true; @@ -128,13 +131,15 @@ TEST(FfiTest, RecordStageCreateAndUpdateCommands) { EXPECT_EQ(ctx.commands().size(), 3); } else if (action == RecordAction::kUpdate) { auto update_launch_st = ctx.UpdateLaunch( - state.launch_cmd, std::initializer_list{}); + /*command=*/state.launch_cmd, + /*args=*/std::initializer_list{}); EXPECT_OK(update_launch_st); state.launch_updated = true; int src = 0, dst = 0; - auto update_memcpy_st = - ctx.UpdateMemcpyD2D(state.memcpy_cmd, &dst, &src, sizeof(int)); + auto update_memcpy_st = ctx.UpdateMemcpyD2D( + /*command=*/state.memcpy_cmd, /*dst=*/&dst, /*src=*/&src, + /*size=*/sizeof(int)); EXPECT_OK(update_memcpy_st); state.memcpy_updated = true; } @@ -150,11 +155,11 @@ TEST(FfiTest, RecordStageCreateAndUpdateCommands) { // Create a record. { XLA_FFI_RecordFrame record_frame = { - reinterpret_cast(&state), - &dummy_record_api, - XLA_FFI_RecordAction_Create, - commands_storage, - &num_commands, + /*record_ctx=*/reinterpret_cast(&state), + /*api=*/&dummy_record_api, + /*action=*/XLA_FFI_RecordAction_Create, + /*commands=*/commands_storage, + /*num_commands=*/&num_commands, /*max_commands=*/10}; RecordExtension::CExtension record_frame_ext = BuildRecordCExtension(&record_frame); @@ -176,11 +181,11 @@ TEST(FfiTest, RecordStageCreateAndUpdateCommands) { // Update the record. { XLA_FFI_RecordFrame record_frame = { - reinterpret_cast(&state), - &dummy_record_api, - XLA_FFI_RecordAction_Update, - commands_storage, - &num_commands, + /*record_ctx=*/reinterpret_cast(&state), + /*api=*/&dummy_record_api, + /*action=*/XLA_FFI_RecordAction_Update, + /*commands=*/commands_storage, + /*num_commands=*/&num_commands, /*max_commands=*/10}; RecordExtension::CExtension record_frame_ext = BuildRecordCExtension(&record_frame); @@ -197,5 +202,106 @@ TEST(FfiTest, RecordStageCreateAndUpdateCommands) { } } +TEST(FfiTest, CreateLaunchWithFunctionPtr) { + struct DummyCommand {}; + DummyCommand launch{}; + const XLA_FFI_Command* expected_cmd = + reinterpret_cast(&launch); + + struct CaptureState { + const XLA_FFI_Command* launch_cmd; + const void* kernel_data = nullptr; + int64_t kernel_size = -1; + XLA_FFI_SourceFormat format = XLA_FFI_SourceFormat_PTX; + } state{expected_cmd}; + + const void* expected_func_ptr = reinterpret_cast(0x12345678); + + XLA_FFI_RecordApi dummy_record_api = {}; + dummy_record_api.create_launch = + +[](XLA_FFI_RecordContext* ctx, const char*, const void* kernel_data, + int64_t kernel_size, XLA_FFI_SourceFormat format, XLA_FFI_LaunchDims, + uint32_t, int32_t, const XLA_FFI_KernelArgs*, + const XLA_FFI_Command* const*, uint32_t, + const XLA_FFI_Command** out_command) -> XLA_FFI_Error* { + auto* s = reinterpret_cast(ctx); + s->kernel_data = kernel_data; + s->kernel_size = kernel_size; + s->format = format; + *out_command = s->launch_cmd; + return nullptr; + }; + + auto handler = Ffi::BindRecord().Ctx>().To( + [&](RecordContext ctx) -> absl::Status { + XLA_FFI_LaunchDims dims{{1, 1, 1}, {1, 1, 1}}; + auto null_cmd = ctx.CreateLaunch( + /*kernel_name=*/"null_func", /*kernel_data=*/nullptr, + /*kernel_size=*/0, /*format=*/SourceFormat::kFunctionPtr, + /*launch_dims=*/dims, /*shared_mem_bytes=*/0, + /*uses_pdl=*/false, /*args=*/std::initializer_list{}); + EXPECT_THAT(null_cmd, + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument)); + + auto launch_cmd = ctx.CreateLaunch( + /*kernel_name=*/"func_ptr_kernel", + /*kernel_data=*/expected_func_ptr, /*kernel_size=*/0, + /*format=*/SourceFormat::kFunctionPtr, /*launch_dims=*/dims, + /*shared_mem_bytes=*/0, /*uses_pdl=*/false, + /*args=*/std::initializer_list{}); + EXPECT_THAT(launch_cmd, IsOkAndHolds(Eq(expected_cmd))); + return absl::OkStatus(); + }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + const XLA_FFI_Command* commands_storage[10] = {nullptr}; + int64_t num_commands = 0; + + XLA_FFI_RecordFrame record_frame = { + /*record_ctx=*/reinterpret_cast(&state), + /*api=*/&dummy_record_api, + /*action=*/XLA_FFI_RecordAction_Create, + /*commands=*/commands_storage, + /*num_commands=*/&num_commands, + /*max_commands=*/10}; + RecordExtension::CExtension record_frame_ext = + BuildRecordCExtension(&record_frame); + + InvokeContext context; + context.extension_start = &record_frame_ext.extension_base; + + ASSERT_OK( + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kRecord)); + EXPECT_EQ(state.kernel_data, expected_func_ptr); + EXPECT_EQ(state.kernel_size, 0); + EXPECT_EQ(state.format, XLA_FFI_SourceFormat_FUNCTION_PTR); + EXPECT_EQ(num_commands, 1); +} + +TEST(FfiTest, RecordExtensionVersionMismatch) { + auto handler = Ffi::BindRecord().Ctx>().To( + [&](RecordContext) -> absl::Status { return absl::OkStatus(); }); + + CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0); + auto call_frame = builder.Build(); + + XLA_FFI_RecordFrame record_frame = {}; + RecordExtension::CExtension record_frame_ext = + BuildRecordCExtension(&record_frame); + record_frame_ext.extension_base.id.minor_version = 1; + + InvokeContext context; + context.extension_start = &record_frame_ext.extension_base; + + auto status = + Invoke(Api(), *handler, call_frame, context, ExecutionStage::kRecord); + EXPECT_FALSE(status.ok()); + EXPECT_THAT( + status.message(), + testing::HasSubstr("Extension version mismatch for RecordExtension")); +} + } // namespace } // namespace xla::ffi diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc index 2f068bbd80f33c..7200b9f496ed82 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc @@ -893,16 +893,17 @@ std::optional PatternMatchParseWhileLoop( // in the type-agnostic handler. For e.g., HandleGetTupleElement in the parent // type-agnostic evaluator will be able to accept Tuple primitive type, whereas // HloEvaluatorTypedVisitor cannot. -HloEvaluator::HloEvaluator(int64_t max_loop_iterations, - bool cache_call_computation_evals, bool is_embedded) +HloEvaluator::HloEvaluator( + int64_t max_loop_iterations, bool cache_call_computation_evals, + std::shared_ptr specialization_cache) : max_loop_iterations_(max_loop_iterations), cache_call_computation_evals_(cache_call_computation_evals), - is_embedded_(is_embedded) { + specialization_cache_(std::move(specialization_cache)) { // Each HandleCall creates its own child evaluator, and in order to have one - // shared cache for all call hierarchy, child evaluators need to borrow caches - // from their parents. Child evaluators has is_embedded true and the - // main/entry one false. - if (cache_call_computation_evals_ && !is_embedded_) { + // shared cache for all call hierarchy, child evaluators borrow the cache + // passed from their parents via CreateEmbedded. If no cache is passed and + // caching is enabled, create a new one. + if (cache_call_computation_evals_ && specialization_cache_ == nullptr) { specialization_cache_ = std::make_shared(); } for (int i = PrimitiveType_MIN; i < PrimitiveType_ARRAYSIZE; ++i) { @@ -967,10 +968,6 @@ absl::StatusOr HloEvaluator::Evaluate( 2, "HloEvaluator::Evaluate computation:\n" + computation.ToString()); OnEvaluateComputation(computation); - if (!is_embedded_) { - ClearSpecializationCache(); - } - if (args.size() != computation.num_parameters()) { return InvalidArgument( "Expected %d argument%s, but got %d.", computation.num_parameters(), @@ -1026,9 +1023,6 @@ absl::StatusOr HloEvaluator::Evaluate( bool recursively_evaluate_nonconstant_operands, const absl::flat_hash_map& substitutions) { - if (!is_embedded_) { - ClearSpecializationCache(); - } ScopedEvaluateState evaluate_state(&state_); // Use the substitutions to manually set instructions results to a specific @@ -3676,9 +3670,7 @@ absl::Status HloEvaluator::HandleCall(const HloInstruction* call) { return absl::OkStatus(); } - if (specialization_cache_ == nullptr) { - specialization_cache_ = std::make_shared(); - } + TF_RET_CHECK(specialization_cache_ != nullptr); const Literal* cached_result = specialization_cache_->Find(computation, arg_literals); diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h index 9f2222ed9d029f..1453825d8a3579 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h @@ -78,19 +78,28 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, const HloDataflowAnalysis* dataflow_analysis; }; + struct SpecializationCache; + // Only evaluate up to max_loop_iterations per while-loop execution if // specified. // - // If `cache_call_computation_evals` true, HloEvaluator caches evals on + // If `cache_call_computation_evals` is true, HloEvaluator caches evals on // computations on call ops to avoid re-evaluating the same call on identical - // constant arguments. + // constant arguments. If `specialization_cache` is non-null, this evaluator + // shares the provided cache; otherwise, if `cache_call_computation_evals` is + // true, a new `SpecializationCache` is created. // - // In case `is_embeeded` is true, HloEvaluator does not create a fresh cache - // for call evals. Instead `CreateEmbedded` api sets it so that the child - // evaluator borrow the cache from the parent evaluator. - explicit HloEvaluator(int64_t max_loop_iterations = -1, - bool cache_call_computation_evals = false, - bool is_embedded = false); + // Note that the cache is keyed by the raw `HloComputation` pointer: + // 1. If a called computation is modified in-place, the cache will not be + // updated or invalidated automatically; users should not enable caching + // if called computations may be modified. + // 2. `HloEvaluator` does not automatically clear the cache between + // evaluations and the user of `HloEvaluator` should call + // `ClearSpecializationCache()` when appropriate. + explicit HloEvaluator( + int64_t max_loop_iterations = -1, + bool cache_call_computation_evals = false, + std::shared_ptr specialization_cache = nullptr); // Returns true if the opcode is implemented by HloEvaluator. False otherwise. static bool IsOpcodeImplemented(HloOpcode opcode); @@ -104,11 +113,10 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, int64_t max_loop_iterations) { auto result = std::make_unique(max_loop_iterations, cache_call_computation_evals_, - /*is_embedded=*/true); + specialization_cache_); result->set_use_fast_path(use_fast_path_); result->set_custom_call_handler(custom_call_handler_); result->set_eval_literal_handler(eval_literal_handler_); - result->specialization_cache_ = specialization_cache_; return result; } @@ -300,7 +308,12 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, // identical constant arguments. // // SpecializationKey identifies an evaluation by its target HloComputation - // and concrete argument literals (compared by value). + // pointer and concrete argument literals. + // + // Because the key uses the raw `HloComputation` pointer, if a computation is + // modified in-place, existing cache entries for that computation pointer are + // not invalidated or updated; callers should not enable caching if + // computations may be modified. // // SpecializationCache stores the mapping from SpecializationKey to the // evaluated result Literal, along with backing storage for argument literals. @@ -407,23 +420,18 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, } }; + // Clears all cached call computation results and argument literals. Must be + // called by the user if computations are modified. void ClearSpecializationCache() { if (specialization_cache_ != nullptr) { specialization_cache_->Clear(); } } - bool is_embedded() const { return is_embedded_; } // Used by unit tests. bool cache_call_computation_evals() const { return cache_call_computation_evals_; } - void set_cache_call_computation_evals(bool enable) { - cache_call_computation_evals_ = enable; - if (cache_call_computation_evals_ && specialization_cache_ == nullptr) { - specialization_cache_ = std::make_shared(); - } - } const SpecializationCache* specialization_cache() const { return specialization_cache_.get(); } @@ -812,8 +820,6 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault, // embedded evaluators to safely share cache ownership without raw-pointer // lifetime issues or manual ownership tracking. std::shared_ptr specialization_cache_; - // True if this evaluator was created as an embedded evaluator. - bool is_embedded_ = false; HloEvaluator(const HloEvaluator&) = delete; HloEvaluator& operator=(const HloEvaluator&) = delete; diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc index 90bb411c9bd494..7b813bacaeeac8 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc @@ -8563,7 +8563,8 @@ TEST_F(HloEvaluatorTest, ConstructorEnablesCacheCallComputationEvals) { } TEST_F(HloEvaluatorTest, MemoizationSameComputationIdenticalArguments) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string = R"( HloModule MemoizationSameArgs @@ -8581,14 +8582,16 @@ TEST_F(HloEvaluatorTest, MemoizationSameComputationIdenticalArguments) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(30.0f), result)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 1); + EXPECT_EQ(evaluator.specialization_cache()->size(), 1); } TEST_F(HloEvaluatorTest, MemoizationDifferentArguments) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string = R"( HloModule MemoizationDiffArgs @@ -8607,14 +8610,16 @@ TEST_F(HloEvaluatorTest, MemoizationDifferentArguments) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(32.0f), result)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 2); + EXPECT_EQ(evaluator.specialization_cache()->size(), 2); } TEST_F(HloEvaluatorTest, MemoizationMultipleParameters) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string = R"( HloModule MemoizationMultiParams @@ -8635,18 +8640,20 @@ TEST_F(HloEvaluatorTest, MemoizationMultipleParameters) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result, + evaluator.Evaluate(*m_->entry_computation(), {})); // call1 = 10 - 3 = 7 // call2 = 10 - 3 = 7 (cache hit) // call3 = 3 - 10 = -7 (cache miss, order matters) // final = 7 + 7 + (-7) = 7 EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(7.0f), result)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 2); + EXPECT_EQ(evaluator.specialization_cache()->size(), 2); } TEST_F(HloEvaluatorTest, MemoizationZeroParameters) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string = R"( HloModule MemoizationZeroParams @@ -8661,14 +8668,16 @@ TEST_F(HloEvaluatorTest, MemoizationZeroParameters) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(84.0f), result)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 1); + EXPECT_EQ(evaluator.specialization_cache()->size(), 1); } TEST_F(HloEvaluatorTest, NestedCallSharedCache) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string = R"( HloModule NestedCalls @@ -8700,7 +8709,8 @@ TEST_F(HloEvaluatorTest, NestedCallSharedCache) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result, + evaluator.Evaluate(*m_->entry_computation(), {})); // comp_b(5) = 15 // comp_a(5) = 15 * 2 = 30 // comp_c(5) = 15 * 3 = 45 @@ -8710,12 +8720,12 @@ TEST_F(HloEvaluatorTest, NestedCallSharedCache) { // comp_a, comp_c, and comp_b should each have 1 entry in the specialization // cache. comp_b with argument 5 is evaluated only once and shared between // comp_a and comp_c. - EXPECT_EQ(evaluator_.specialization_cache()->size(), 3); + EXPECT_EQ(evaluator.specialization_cache()->size(), 3); const HloComputation* comp_b = m_->GetComputationWithName("comp_b"); ASSERT_NE(comp_b, nullptr); int comp_b_entries = 0; - for (const auto& [key, val] : *evaluator_.specialization_cache()) { + for (const auto& [key, val] : *evaluator.specialization_cache()) { if (key.computation == comp_b) { ++comp_b_entries; EXPECT_TRUE( @@ -8726,7 +8736,8 @@ TEST_F(HloEvaluatorTest, NestedCallSharedCache) { } TEST_F(HloEvaluatorTest, ClearSpecializationCache) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string = R"( HloModule ClearCacheModule @@ -8742,24 +8753,27 @@ TEST_F(HloEvaluatorTest, ClearSpecializationCache) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), result)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 1); + EXPECT_EQ(evaluator.specialization_cache()->size(), 1); - evaluator_.ClearSpecializationCache(); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 0); + evaluator.ClearSpecializationCache(); + EXPECT_EQ(evaluator.specialization_cache()->size(), 0); - // Re-evaluating on the top-level evaluator should repopulate the cache - evaluator_.ResetVisitStates(); - ASSERT_OK_AND_ASSIGN(Literal result2, Evaluate()); + // Re-evaluating on the evaluator should repopulate the cache + evaluator.ResetVisitStates(); + ASSERT_OK_AND_ASSIGN(Literal result2, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), result2)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 1); + EXPECT_EQ(evaluator.specialization_cache()->size(), 1); } TEST_F(HloEvaluatorTest, MemoizationConsecutiveTopLevelEvaluations) { - evaluator_.set_cache_call_computation_evals(true); + HloEvaluator evaluator(/*max_loop_iterations=*/-1, + /*cache_call_computation_evals=*/true); const char* const hlo_string1 = R"( HloModule Module1 @@ -8775,10 +8789,14 @@ TEST_F(HloEvaluatorTest, MemoizationConsecutiveTopLevelEvaluations) { } )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string1)); - ASSERT_OK_AND_ASSIGN(Literal result1, Evaluate()); + ASSERT_OK_AND_ASSIGN(Literal result1, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(15.0f), result1)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 1); + EXPECT_EQ(evaluator.specialization_cache()->size(), 1); + + // Explicitly clearing the cache removes all entries. + evaluator.ClearSpecializationCache(); const char* const hlo_string2 = R"( HloModule Module2 @@ -8796,13 +8814,16 @@ TEST_F(HloEvaluatorTest, MemoizationConsecutiveTopLevelEvaluations) { )"; ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_string2)); // Reset visitor state before running visitor on the new module. - evaluator_.ResetVisitStates(); - // Second top-level Evaluate() must automatically clear previous cache - // entries. - ASSERT_OK_AND_ASSIGN(Literal result2, Evaluate()); + evaluator.ResetVisitStates(); + ASSERT_OK_AND_ASSIGN(Literal result2, + evaluator.Evaluate(*m_->entry_computation(), {})); EXPECT_TRUE( LiteralTestUtil::Equal(LiteralUtil::CreateR0(27.0f), result2)); - EXPECT_EQ(evaluator_.specialization_cache()->size(), 1); + EXPECT_EQ(evaluator.specialization_cache()->size(), 1); + + // Explicitly clearing the cache removes all entries. + evaluator.ClearSpecializationCache(); + EXPECT_EQ(evaluator.specialization_cache()->size(), 0); } TEST(EvalErrorTest, OK) { diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index 33513636280a9c..aa10a9def8c674 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -307,6 +307,7 @@ class HloInstruction { inline static constexpr char kMainExecutionThread[] = "main"; inline static constexpr char kHostThread[] = "host"; + inline static constexpr char kParallelExecutionThread[] = "parallel"; // Iota based id unique inside parent computation. using LocalId = int32_t; diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h index 33d5a24cf31e64..da532efae24e83 100644 --- a/third_party/xla/xla/pjrt/transpose_kernels.h +++ b/third_party/xla/xla/pjrt/transpose_kernels.h @@ -146,6 +146,60 @@ HWY_INLINE void UnpackStepCPS(Cont&& cont, Vs&... in) { std::forward(cont), std::make_index_sequence{}, in...); } +// CPS: InterleaveWithinBlocks (pack -> pack) +template +HWY_INLINE void InterleaveWithinBlocksCPS(Cont&& cont, Vs... in) { + if constexpr (element_size * step_size < unpack_limit && + element_size * step_size <= kMaxLaneBytes) { + UnpackStepCPS( + [&](auto... out) { + // out... are by-value parameters; do not forward as rvalues. + InterleaveWithinBlocksCPS( + std::forward(cont), out...); + }, + in...); // lvalues inside this function + } else { + std::forward(cont)(in...); + } +} + +// Helper template that applies the pairwise interleaving reduction. +template +HWY_INLINE void CombinePairsAccImpl(Cont&& cont, std::index_sequence, + Vs&&... vs) { + std::forward(cont)(Unpack( + GetFromPack<2 * I>(std::forward(vs)...), + GetFromPack<2 * I + 1>(std::forward(vs)...))...); +} + +// Applies a pairwise interleaving reduction across an arbitrary sequence of +// vectors. Adjacent vectors are combined using their lower halves, and the +// resulting halved-length sequence is forwarded to the provided continuation. +template +HWY_INLINE void CombinePairsAcc(Cont&& cont, Vs&&... in) { + CombinePairsAccImpl( + std::forward(cont), std::make_index_sequence{}, + std::forward(in)...); +} + +template +HWY_INLINE void CombineRowsCPS(Cont&& cont, Vs... in) { + constexpr size_t N = sizeof...(Vs); + if constexpr (N > 1 && element_size * bs < vector_bytes) { + static_assert(N % 2 == 0); + CombinePairsAcc( + [&](auto... paired) { + CombineRowsCPS( + std::forward(cont), paired...); + }, + in...); + } else { + std::forward(cont)(in...); + } +} + } // namespace HWY_NAMESPACE } // namespace xla HWY_AFTER_NAMESPACE(); diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_test.cc index b3e6757086d134..00ae3fb4a280c7 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler_test.cc @@ -99,6 +99,7 @@ ENTRY main { DebugOptions debug_options = GetDebugOptionsForTest(); debug_options.add_xla_cpu_experimental_ynn_fusion_type( xla::DebugOptions::LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT); + debug_options.set_xla_cpu_experimental_onednn_custom_call(false); debug_options.mutable_xla_backend_extra_options()->insert( {"xla_is_host_offload", "true"}); config.set_debug_options(debug_options); diff --git a/third_party/xla/xla/service/cpu/onednn_util.h b/third_party/xla/xla/service/cpu/onednn_util.h index 9d09b00224f127..2aa9d51778c2d8 100644 --- a/third_party/xla/xla/service/cpu/onednn_util.h +++ b/third_party/xla/xla/service/cpu/onednn_util.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/log/log.h" #include "absl/status/statusor.h" #include "unsupported/Eigen/CXX11/Tensor" #include "oneapi/dnnl/dnnl.hpp" @@ -51,6 +52,19 @@ inline bool IsSupportedType(xla::PrimitiveType dtype) { (TestCPUFeature(CPUFeature::AVX512_FP16) || TestCPUFeature(CPUFeature::AMX_FP16))) || TestCPUFeature(CPUFeature::AVX_NE_CONVERT); + case F8E5M2: + case F8E4M3FN: + case F8E4M3: + if (TestCPUFeature(CPUFeature::AMX_FP8)) { + return true; + } + if (TestCPUFeature(CPUFeature::AVX512BW) && + TestCPUFeature(CPUFeature::AMX_FP16)) { + LOG_FIRST_N(INFO, 1) << "XLA:CPU FP8 dispatched via oneDNN AMX-FP16 " + "emulation path (no native AMX-FP8 on host)."; + return true; + } + return false; default: return false; } diff --git a/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc index f1f82e25aa7864..f6a9a862a24456 100644 --- a/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc +++ b/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc @@ -214,6 +214,26 @@ std::vector GetX86MaxIsaTestCases() { MaxIsaTestSpec{"AVX512", "avx512f", true}, MaxIsaTestSpec{"AVX512", "avx512vnni", false}, MaxIsaTestSpec{"AVX512", "amx-bf16", false}, + MaxIsaTestSpec{"AVX512", "amx-fp8", false}, + // AMX_FP16 as max: SPR-or-older ISAs stay on; DMR additions (amx-fp8, + // other amx tiles, avx10, APX) are suppressed. + MaxIsaTestSpec{"AMX_FP16", "amx-bf16", true}, + MaxIsaTestSpec{"AMX_FP16", "amx-int8", true}, + MaxIsaTestSpec{"AMX_FP16", "avx512bf16", true}, + MaxIsaTestSpec{"AMX_FP16", "amx-fp16", true}, + MaxIsaTestSpec{"AMX_FP16", "avx10.1", true}, + MaxIsaTestSpec{"AMX_FP16", "amx-fp8", false}, + MaxIsaTestSpec{"AMX_FP16", "amx-tf32", false}, + MaxIsaTestSpec{"AMX_FP16", "amx-avx512", false}, + MaxIsaTestSpec{"AMX_FP16", "avx10.2", false}, + MaxIsaTestSpec{"AMX_FP16", "ndd", false}, + MaxIsaTestSpec{"AMX_FP16", "ppx", false}, + // AMX_FP8 as max: amx-fp8 and older AMX tiles stay on. + MaxIsaTestSpec{"AMX_FP8", "amx-fp8", true}, + MaxIsaTestSpec{"AMX_FP8", "amx-fp16", true}, + MaxIsaTestSpec{"AMX_FP8", "amx-bf16", true}, + // Older max still suppresses amx-fp8. + MaxIsaTestSpec{"AMX", "amx-fp8", false}, }); } diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 557efae12c9a3e..8fdb54cd4a0841 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -2867,7 +2867,6 @@ cc_library( "//xla:xla_proto_cc", "//xla/backends/gpu/transforms:pgle_accuracy_checker", "//xla/backends/gpu/transforms:scheduling_instruction_annotator", - "//xla/backends/gpu/transforms:stream_attribute_async_wrapper", "//xla/backends/gpu/transforms/collectives:async_collective_annotator", "//xla/backends/gpu/transforms/collectives:collective_domain", "//xla/backends/gpu/transforms/collectives:collective_ops_utils", diff --git a/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h b/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h index f08a3f4555f162..68c8bdad8e7437 100644 --- a/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h +++ b/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h @@ -34,7 +34,7 @@ class AutotuneCacheKey { // changes that may affect the autotuning results. // To prevent accidental merges of concurrent increments, update the comment // to explain why the version is bumped. - static constexpr int kCurrentVersion = 52; // Triton integration 1.38 + static constexpr int kCurrentVersion = 53; // New block-tiling is a default. AutotuneCacheKey(const se::DeviceDescription& device_description, const HloInstruction& instruction, diff --git a/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc b/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc index 1a56e8474c1ecd..ab225d333a7762 100644 --- a/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc +++ b/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc @@ -45,7 +45,6 @@ limitations under the License. #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h" #include "xla/backends/gpu/transforms/pgle_accuracy_checker.h" #include "xla/backends/gpu/transforms/scheduling_instruction_annotator.h" -#include "xla/backends/gpu/transforms/stream_attribute_async_wrapper.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_input_output_alias_config.h" @@ -854,7 +853,7 @@ absl::StatusOr ScheduleGpuModuleWithMemoryScheduler( PostProcessSchedule), /*execution_threads=*/ {HloInstruction::kMainExecutionThread, - StreamAttributeAsyncWrapper::kParallelExecutionThread}, + HloInstruction::kParallelExecutionThread}, peak_memory_bytes); } diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD index 6c229543e76237..96be346dcc1781 100644 --- a/third_party/xla/xla/service/gpu/model/BUILD +++ b/third_party/xla/xla/service/gpu/model/BUILD @@ -628,6 +628,7 @@ xla_cc_test( "//xla/codegen/tiling:tiling_specification", "//xla/codegen/tiling/experimental:tile", "//xla/codegen/tiling/experimental:tiled_hlo", + "//xla/codegen/tiling/experimental:tiling_space_utils", "//xla/codegen/xtile:block_level_parameters", "//xla/codegen/xtile/codegen:emitter_helpers", "//xla/hlo/analysis:symbolic_map", @@ -635,7 +636,6 @@ xla_cc_test( "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", - "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_device_info_for_tests", "//xla/service/gpu:ir_emission_utils", "//xla/stream_executor:device_description", diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc index 5327b8289cdaaa..0354c406cde3ea 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc @@ -29,6 +29,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/codegen/tiling/experimental/tiled_hlo.h" #include "xla/codegen/tiling/experimental/tiling_space.h" +#include "xla/codegen/tiling/experimental/tiling_space_utils.h" #include "xla/codegen/tiling/symbolic_tile_analysis.h" #include "xla/codegen/tiling/tiled_hlo_computation.h" #include "xla/codegen/tiling/tiling_specification.h" @@ -41,7 +42,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/utils/hlo_traversal.h" -#include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" diff --git a/third_party/xla/xla/service/memory_space_assignment/BUILD b/third_party/xla/xla/service/memory_space_assignment/BUILD index 2deeb91287fc44..88f1ca2cd367fe 100644 --- a/third_party/xla/xla/service/memory_space_assignment/BUILD +++ b/third_party/xla/xla/service/memory_space_assignment/BUILD @@ -1,17 +1,7 @@ # Description: -# Memory Space Assignment service implementation. +# Pre-set assignments data structure for BufferAssignment. -load( - "//xla:xla.default.bzl", - "xla_cc_test", - "xla_py_proto_library", -) -load("//xla/tests:build_defs.bzl", "xla_test") load("//xla/tsl:tsl.bzl", "internal_visibility") -load( - "//xla/tsl/platform:build_config.bzl", - "tf_proto_library", -) load("//xla/tsl/platform:rules_cc.bzl", "cc_library") package( @@ -27,734 +17,19 @@ package_group( ], ) -tf_proto_library( - name = "memory_space_assignment_proto", - srcs = ["memory_space_assignment.proto"], - make_default_target_header_only = True, -) - -xla_py_proto_library( - name = "memory_space_assignment_proto_py_pb2", - deps = [":memory_space_assignment_proto"], -) - cc_library( name = "memory_space_assignment", - srcs = ["memory_space_assignment.cc"], hdrs = ["memory_space_assignment.h"], deps = [ - ":algorithm", - ":allocation", - ":memory_space_assignment_proto_cc", - ":options", - ":simulator", - ":slice", - "//xla:shape_util", - "//xla:status_macros", - "//xla:util", - "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_buffer", "//xla/service:hlo_proto_cc", "//xla/service:hlo_value", "//xla/service/heap_simulator", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/base", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/log:vlog_is_on", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:str_format", - "@com_google_absl//absl/types:span", - "@tsl//tsl/platform:casts", - ], -) - -xla_test( - name = "memory_space_assignment_test", - srcs = ["memory_space_assignment_test.cc"], - backends = [ - "cpu", - ], - deps = [ - ":algorithm", - ":allocation", - ":allocation_value", - ":best_fit_repacker", - ":buffer_interval_comparator", - ":cost_analysis", - ":memory_space_assignment", - ":memory_space_assignment_proto_cc", - ":memory_space_assignment_test_base", - ":options", - ":prefetch_interval_picker", - ":repacking", - ":slice", - ":testing_utils", - ":utils", - "//xla:comparison_util", - "//xla:literal_util", - "//xla:shape_util", - "//xla:util", - "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/analysis:hlo_dataflow_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:verified_hlo_module", - "//xla/hlo/transforms/simplifiers:instruction_hoister", - "//xla/hlo/utils:hlo_live_range", - "//xla/hlo/utils:hlo_matchers", - "//xla/service:hlo_buffer", - "//xla/service:hlo_cost_analysis", - "//xla/service:hlo_proto_cc", - "//xla/service:hlo_value", - "//xla/service/cost_modelling:op_cost", - "//xla/service/heap_simulator", - "//xla/service/heap_simulator:allocation_block", - "//xla/tests:test_utils", - "//xla/tests:xla_internal_test_main", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:logging", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:status_matchers", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - "@com_google_googletest//:gtest", - "@tsl//tsl/platform:protobuf", - ], -) - -cc_library( - name = "repacking", - hdrs = ["repacking.h"], - deps = [ - "//xla/service/heap_simulator:allocation_block", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", - ], -) - -cc_library( - name = "best_fit_repacker", - srcs = ["best_fit_repacker.cc"], - hdrs = ["best_fit_repacker.h"], - deps = [ - ":repacking", - "//xla:comparison_util", - "//xla/service/heap_simulator", - "//xla/service/heap_simulator:allocation_block", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/functional:any_invocable", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/log:vlog_is_on", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], ) cc_library( - name = "memory_space_assignment_test_base", - testonly = True, - hdrs = ["memory_space_assignment_test_base.h"], - deps = [ - ":buffer_interval_comparator", - ":cost_analysis", - ":memory_space_assignment", - ":options", - ":prefetch_interval_picker", - ":utils", - "//xla:shape_util", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/transforms/simplifiers:instruction_hoister", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:buffer_value", - "//xla/service:hlo_buffer", - "//xla/service:hlo_cost_analysis", - "//xla/service:hlo_value", - "//xla/service/cost_modelling:op_cost", - "//xla/tests:hlo_pjrt_test_base", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", - "@com_google_googletest//:gtest_for_library", - ], -) - -cc_library( - name = "utils", - srcs = ["utils.cc"], - hdrs = ["utils.h"], - deps = [ - ":memory_space_assignment_proto_cc", - "//xla:shape_util", - "//xla:util", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_value", - "//xla/service/heap_simulator", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", - "@com_googlesource_code_re2//:re2", - "@highwayhash", - "@highwayhash//:arch_specific", - "@highwayhash//:hh_types", - ], -) - -cc_library( - name = "slice", - srcs = ["slice.cc"], - hdrs = [ - "slice.h", - ], - deps = [ - ":memory_space_assignment_proto_cc", - "//xla:shape_util", - "//xla/service:time_utils", - "//xla/service/heap_simulator", - "@com_google_absl//absl/functional:any_invocable", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - ], -) - -xla_cc_test( - name = "slice_test", - srcs = ["slice_test.cc"], - deps = [ - ":slice", - "//xla/service:time_utils", - "@com_google_absl//absl/functional:any_invocable", - "@com_google_absl//absl/log:check", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "allocation", - srcs = ["allocation.cc"], - hdrs = [ - "allocation.h", - ], - deps = [ - ":memory_space_assignment_proto_cc", - ":slice", - "//xla:shape_util", - "//xla:status_macros", - "//xla:util", - "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_value", - "//xla/service:time_utils", - "//xla/service:tuple_util", - "//xla/service/heap_simulator", - "//xla/service/heap_simulator:allocation_block", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/base", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/functional:function_ref", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - "@tsl//tsl/platform:casts", - "@tsl//tsl/platform:errors", - "@tsl//tsl/platform:statusor", - ], -) - -xla_cc_test( - name = "allocation_test", - srcs = ["allocation_test.cc"], - deps = [ - ":allocation", - "//xla:shape_util", - "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_value", - "//xla/service/heap_simulator", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "tuning_utils", - srcs = ["tuning_utils.cc"], - hdrs = ["tuning_utils.h"], - deps = [ - "//xla/hlo/ir:hlo", - "//xla/service:hlo_value", - "//xla/service/heap_simulator", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/log:check", - ], -) - -cc_library( - name = "options", - srcs = ["options.cc"], - hdrs = ["options.h"], - deps = [ - ":allocation_value", - ":buffer_interval_comparator", - ":cost_analysis", - ":memory_space_assignment_proto_cc", - ":prefetch_interval_picker", - ":repacking", - ":slice", - "//xla:shape_util", - "//xla:util", - "//xla/hlo/ir:hlo", - "//xla/service:buffer_value", - "//xla/service:hlo_value", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", - "@com_google_absl//absl/types:span", - ], -) - -cc_library( - name = "cost_analysis", - srcs = ["cost_analysis.cc"], - hdrs = ["cost_analysis.h"], - deps = [ - ":utils", - "//xla:shape_util", - "//xla:util", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/analysis:while_loop_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:call_graph", - "//xla/service:hlo_buffer", - "//xla/service:hlo_value", - "//xla/service/cost_modelling:op_cost", - "//xla/service/heap_simulator", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/functional:function_ref", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings:string_view", - "@com_google_absl//absl/types:span", - ], -) - -xla_cc_test( - name = "cost_analysis_test", - srcs = ["cost_analysis_test.cc"], - deps = [ - ":cost_analysis", - "//xla:shape_util", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/service:hlo_cost_analysis", - "//xla/service/cost_modelling:op_cost", - "@com_google_absl//absl/log", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:status_matchers", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - "@tsl//tsl/platform:errors", - ], -) - -cc_library( - name = "simulator", - srcs = ["simulator.cc"], - hdrs = ["simulator.h"], - deps = [ - ":allocation", - ":cost_analysis", - "//xla:shape_util", - "//xla:util", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - ], -) - -xla_cc_test( - name = "simulator_test", - srcs = ["simulator_test.cc"], - deps = [ - ":allocation", - ":cost_analysis", - ":simulator", - "//xla:shape_util", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_cost_analysis", - "//xla/service:hlo_value", - "//xla/service/cost_modelling:op_cost", - "//xla/service/heap_simulator", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - "@tsl//tsl/platform:errors", - ], -) - -cc_library( - name = "prefetch_interval_picker", - srcs = ["prefetch_interval_picker.cc"], - hdrs = ["prefetch_interval_picker.h"], - deps = [ - ":cost_analysis", - ":memory_space_assignment_proto_cc", - "//xla:shape_util", - "//xla:util", - "//xla:xla_data_proto_cc", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_proto_cc", - "//xla/service:hlo_value", - "//xla/service/heap_simulator", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - "@tsl//tsl/platform:logging", - ], -) - -cc_library( - name = "testing_utils", - testonly = True, - hdrs = ["testing_utils.h"], - deps = [ - ":cost_analysis", - "//xla:shape_util", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:call_graph", - "//xla/service:hlo_cost_analysis", - "//xla/service/cost_modelling:op_cost", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", - "@tsl//tsl/platform:statusor", - ], -) - -cc_library( - name = "memory_bound_loop_optimizer", - srcs = ["memory_bound_loop_optimizer.cc"], - hdrs = ["memory_bound_loop_optimizer.h"], - deps = [ - ":allocation", - ":cost_analysis", - ":memory_space_assignment_proto_cc", - ":options", - "//xla:shape_util", - "//xla:status_macros", - "//xla:util", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:buffer_value", - "//xla/service:hlo_buffer", - "//xla/service:hlo_proto_cc", - "//xla/service:hlo_value", - "//xla/service/heap_simulator", - "//xla/service/heap_simulator:allocation_block", - "//xla/tsl/platform:errors", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/base:nullability", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/log:die_if_null", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:str_format", - "@com_google_absl//absl/types:span", - ], -) - -xla_cc_test( - name = "memory_bound_loop_optimizer_test", - srcs = ["memory_bound_loop_optimizer_test.cc"], - deps = [ - ":allocation", - ":buffer_interval_comparator", - ":cost_analysis", - ":memory_bound_loop_optimizer", - ":memory_space_assignment", - ":memory_space_assignment_proto_cc", - ":options", - ":prefetch_interval_picker", - "//xla:shape_util", - "//xla:status_macros", - "//xla:util", - "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/hlo/testlib:verified_hlo_module", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:buffer_value", - "//xla/service:hlo_cost_analysis", - "//xla/service:hlo_value", - "//xla/service/cost_modelling:op_cost", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - "@com_googlesource_code_re2//:re2", - ], -) - -cc_library( - name = "live_range_util", - srcs = ["live_range_util.cc"], - hdrs = ["live_range_util.h"], - deps = [ - "//xla:util", - "//xla/hlo/ir:hlo", - "//xla/service:hlo_buffer", - "//xla/service:hlo_value", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/strings", - ], -) - -xla_cc_test( - name = "live_range_util_test", - srcs = ["live_range_util_test.cc"], - deps = [ - ":live_range_util", - "//xla:shape_util", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:hlo_buffer", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "algorithm", - srcs = ["algorithm.cc"], - hdrs = ["algorithm.h"], - deps = [ - ":allocation", - ":allocation_value", - ":buffer_interval_comparator", - ":cost_analysis", - ":live_range_util", - ":memory_bound_loop_optimizer", - ":memory_space_assignment_proto_cc", - ":options", - ":repacking", - ":slice", - ":tuning_utils", - ":utils", - "//xla:debug_options_flags", - "//xla:frontend_attributes", - "//xla:shape_tree", - "//xla:shape_util", - "//xla:status_macros", - "//xla:util", - "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_alias_analysis", - "//xla/hlo/analysis:hlo_dataflow_analysis", - "//xla/hlo/analysis:hlo_operand_index", - "//xla/hlo/ir:hlo", - "//xla/hlo/ir:hlo_instruction_utils", - "//xla/hlo/transforms:memory_space_propagation", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:buffer_value", - "//xla/service:call_graph", - "//xla/service:computation_layout", - "//xla/service:decision", - "//xla/service:hlo_buffer", - "//xla/service:hlo_proto_cc", - "//xla/service:hlo_value", - "//xla/service:time_utils", - "//xla/service/heap_simulator", - "//xla/service/heap_simulator:allocation_block", - "//xla/tsl/platform:logging", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:btree", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/container:inlined_vector", - "@com_google_absl//absl/container:linked_hash_map", - "@com_google_absl//absl/functional:any_invocable", - "@com_google_absl//absl/hash", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/log:vlog_is_on", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:str_format", - "@com_google_absl//absl/strings:string_view", - "@com_google_absl//absl/types:span", - ], -) - -cc_library( - name = "buffer_interval_comparator", - srcs = ["buffer_interval_comparator.cc"], - hdrs = ["buffer_interval_comparator.h"], - deps = [ - ":cost_analysis", - ":memory_space_assignment_proto_cc", - ":utils", - "//xla:shape_util", - "//xla:util", - "//xla/hlo/ir:hlo", - "//xla/hlo/utils:hlo_live_range", - "//xla/service:buffer_value", - "//xla/service:hlo_value", - "//xla/service/heap_simulator", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", - ], -) - -cc_library( - name = "allocation_value", - srcs = ["allocation_value.cc"], - hdrs = ["allocation_value.h"], - deps = [ - ":allocation", - "//xla:shape_util", - "//xla/hlo/ir:hlo", - "//xla/service:hlo_value", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - ], -) - -xla_cc_test( - name = "prefetch_interval_picker_test", - srcs = ["prefetch_interval_picker_test.cc"], - deps = [ - ":cost_analysis", - ":prefetch_interval_picker", - ":testing_utils", - "//xla:shape_util", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/service:hlo_cost_analysis", - "//xla/service:hlo_value", - "//xla/service/cost_modelling:op_cost", - "@com_google_absl//absl/log", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - ], -) - -xla_cc_test( - name = "best_fit_repacker_test", - srcs = ["best_fit_repacker_test.cc"], - deps = [ - ":best_fit_repacker", - "//xla:comparison_util", - "//xla/service/heap_simulator", - "//xla/service/heap_simulator:allocation_block", - "//xla/tests:xla_internal_test_main", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/types:span", - "@com_google_googletest//:gtest", - "@tsl//tsl/platform:test", - ], + name = "memory_space_assignment_proto_cc_impl", + visibility = ["//visibility:public"], ) diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc deleted file mode 100644 index f765dad0640671..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ /dev/null @@ -1,11102 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/algorithm.h" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/container/inlined_vector.h" -#include "absl/container/linked_hash_map.h" -#include "absl/functional/any_invocable.h" -#include "absl/hash/hash.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/log/vlog_is_on.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_format.h" -#include "absl/strings/str_join.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/debug_options_flags.h" -#include "xla/frontend_attributes.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/analysis/hlo_dataflow_analysis.h" -#include "xla/hlo/analysis/hlo_operand_index.h" -#include "xla/hlo/ir/hlo_casting_utils.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_input_output_alias_config.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_instruction_utils.h" -#include "xla/hlo/ir/hlo_instructions.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/ir/hlo_schedule.h" -#include "xla/hlo/transforms/memory_space_propagation.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/layout.h" -#include "xla/map_util.h" -#include "xla/service/buffer_value.h" -#include "xla/service/call_graph.h" -#include "xla/service/computation_layout.h" -#include "xla/service/decision.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/allocation_value.h" -#include "xla/service/memory_space_assignment/buffer_interval_comparator.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/live_range_util.h" -#include "xla/service/memory_space_assignment/memory_bound_loop_optimizer.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/service/memory_space_assignment/repacking.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/service/memory_space_assignment/tuning_utils.h" -#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/status_macros.h" -#include "xla/tsl/platform/logging.h" -#include "xla/util.h" -#include "xla/xla_data.pb.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -bool IsSliceLikeInstruction(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kSlice || - instruction->opcode() == HloOpcode::kDynamicSlice; -} - -// Returns true if the instruction is a trivial instruction. -bool IsTrivialInstruction(const HloInstruction* inst) { - return inst->opcode() == HloOpcode::kGetTupleElement || - inst->opcode() == HloOpcode::kTuple || - inst->opcode() == HloOpcode::kBitcast; -} - -// Returns true if the position is a trivial position. -bool IsTrivialPosition(HloPosition position) { - return IsTrivialInstruction(position.instruction); -} - -// All trivial uses that are non root do not require a separate allocation. -// - They are are no ops, freely movable and we can even discard them and add -// them back later. -// - We do not need to extend the allocation for these uses, so we can just skip -// them. -// If they are a root instruction we need to extend the allocation till the root -// instruction so the uses outside the computation can use it. -// Returns true if the use is a trivial but non-root instruction. -bool IsUseTrivialNonRoot(HloUse use) { - return IsTrivialInstruction(use.instruction) && - use.instruction != use.instruction->parent()->root_instruction(); -} - -HloPosition GetNonTrivialSourcePosition(HloPosition position) { - while (IsTrivialPosition(position)) { - if (position.instruction->opcode() == HloOpcode::kGetTupleElement) { - int64_t tuple_index = position.instruction->tuple_index(); - position.instruction = position.instruction->mutable_operand(0); - position.index.push_front(tuple_index); - } else if (position.instruction->opcode() == HloOpcode::kTuple) { - if (position.index.empty()) { - return position; - } - int64_t tuple_index = position.index.front(); - position.index.pop_front(); - position.instruction = position.instruction->mutable_operand(tuple_index); - } else if (position.instruction->opcode() == HloOpcode::kBitcast) { - position.instruction = position.instruction->mutable_operand(0); - } - } - return position; -} - -bool IsAsyncOperationStateDefinition(const HloInstruction* instruction) { - return HloDataflowAnalysis::IsAsynchronousOperationStart( - instruction->opcode()) || - instruction->opcode() == HloOpcode::kAsyncUpdate; -} - -bool IsAsyncOperationStateUse(const HloUse& use) { - return use.operand_number == 0 && - (HloDataflowAnalysis::IsAsynchronousOperationDone( - use.instruction->opcode()) || - use.instruction->opcode() == HloOpcode::kAsyncUpdate); -} - -// Define a dummy chunk for chunks that will be allocated in the default memory -// space and for keeping track of number of asynchronous copies. -const HeapSimulator::Chunk kDummyChunk = - HeapSimulator::Chunk::FromOffsetSize(-1, -1); -// For cross-program prefetched buffer, we only perform the freeing optimization -// if the buffer occupies less of the execution time ratio than this value. -const float kCrossProgramPrefetchOccupyFreeingLimit = 0.6; - -int64_t GetAlignedOffset(int64_t offset, int64_t alignment) { - return CeilOfRatio(offset, alignment) * alignment; -} - -// Returns true if 'instruction' is a while loop configured with disabled loop -// copies (e.g. for async while loop pipelining). -bool IsAsyncPipelinedWhileLoop(const HloInstruction* instruction) { - return instruction != nullptr && instruction->opcode() == HloOpcode::kWhile && - HasDisableWhileLoopCopiesAttr(instruction); -} - -// Returns true if the module contains any while loops configured with disabled -// loop copies (e.g. for async while loop pipelining). -bool HasAsyncPipelinedWhileLoops(const HloModule& module) { - for (const HloComputation* computation : module.computations()) { - for (const HloInstruction* instruction : computation->instructions()) { - if (IsAsyncPipelinedWhileLoop(instruction)) { - return true; - } - } - } - return false; -} - -// Returns true if 'buffer' is defined by, used by, or contained within an -// async pipelined while loop computation. -bool IsBufferAliasedToAsyncPipelinedWhileLoop(const HloValue* buffer) { - if (buffer == nullptr) { - return false; - } - for (const HloPosition& pos : buffer->positions()) { - if (IsAsyncPipelinedWhileLoop(pos.instruction)) { - return true; - } - if (pos.instruction->parent() != nullptr) { - for (const HloInstruction* caller : - pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { - if (IsAsyncPipelinedWhileLoop(caller)) { - return true; - } - } - } - } - for (const HloUse& use : buffer->GetUses()) { - if (IsAsyncPipelinedWhileLoop(use.instruction)) { - return true; - } - if (use.instruction->parent() != nullptr) { - for (const HloInstruction* caller : - use.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { - if (IsAsyncPipelinedWhileLoop(caller)) { - return true; - } - } - } - } - return false; -} - -template -std::string VectorToString(const std::vector& v, - bool include_indices = false, int start = 0, - int end = std::numeric_limits::max()) { - std::vector elements; - - for (int i = start; i < end && i < v.size(); ++i) { - std::string prefix; - if (include_indices) { - prefix = absl::StrCat(i, ": "); - } - elements.push_back(absl::StrCat(prefix, v[i])); - } - - return absl::StrCat("[ ", absl::StrJoin(elements, ", "), " ]"); -} - -bool LooksLikeAnActivation(const HloInstruction* inst, bool permissive_mode) { - for (HloInstruction* user : inst->users()) { - switch (user->opcode()) { - case HloOpcode::kConvolution: - case HloOpcode::kDot: - if (user->operand(0) == inst) { - return true; - } - break; - case HloOpcode::kGather: - if (user->operand(1) == inst) { - return true; - } - break; - case HloOpcode::kFusion: - for (int i = 0; i < user->operand_count(); ++i) { - if (user->operand(i) == inst && - LooksLikeAnActivation(user->fused_parameter(i), - permissive_mode)) { - return true; - } - } - break; - case HloOpcode::kBitcast: - case HloOpcode::kBroadcast: - case HloOpcode::kTranspose: - if (LooksLikeAnActivation(user, permissive_mode)) { - return true; - } - break; - case HloOpcode::kCopy: - if (user->IsFused() && (user == user->parent()->root_instruction())) { - user = user->parent()->FusionInstruction(); - if (LooksLikeAnActivation(user, permissive_mode)) { - return true; - } - break; - } - return true; - case HloOpcode::kDynamicUpdateSlice: - case HloOpcode::kDynamicSlice: - if (std::find(user->operands().begin() + 1, user->operands().end(), - inst) != user->operands().end()) { - return true; - } - if (LooksLikeAnActivation(user, permissive_mode)) { - return true; - } - break; - case HloOpcode::kReduce: - // Check init operands. - if (std::find(user->operands().begin() + user->operand_count() / 2, - user->operands().end(), inst) != user->operands().end()) { - return true; - } - if (LooksLikeAnActivation(user, permissive_mode)) { - return true; - } - break; - default: - // Permissive mode assumes the tensor is not an activation when we - // couldn't explicitly determine that it is not an activation. - return !permissive_mode; - } - } - return false; -} - -// Returns true if the use value does not live out of the module. The value -// lives out if it is the root or it aliases with another value that lives out. -// We recurse to detect the latter case. -bool UseDoesNotLiveOut( - const HloUse& use, const HloAliasAnalysis& alias_analysis, - const AliasInfo* alias_info, const HloInstruction* root_instruction, - absl::flat_hash_map& use_to_does_not_live_out) { - if (use.instruction == root_instruction && - (use.instruction->opcode() == HloOpcode::kTuple || - use.instruction->opcode() == HloOpcode::kBitcast)) { - return false; - } - // If already evaluated or on the current recursion stack, return the result. - auto it = use_to_does_not_live_out.find(use); - if (it != use_to_does_not_live_out.end()) { - return it->second; - } - // Mark this use as in-progress (defaulting to true) to break in-place - // aliasing cycles (e.g. async start <-> done). - use_to_does_not_live_out[use] = true; - auto in_place_pairs = alias_info->GetInPlaceInputOutputPairs(use.instruction); - bool does_not_live_out = absl::c_all_of( - in_place_pairs, - [&](const std::pair& in_place_pair) { - if (in_place_pair.first.operand_number == use.operand_number && - in_place_pair.first.operand_index == use.operand_index) { - if (use.instruction == root_instruction) { - return false; - } - for (const HloUse& nested_use : - alias_analysis.dataflow_analysis() - .GetUniqueValueAt(use.instruction, in_place_pair.second) - .GetUses()) { - if (nested_use != use && - !UseDoesNotLiveOut(nested_use, alias_analysis, alias_info, - root_instruction, - use_to_does_not_live_out)) { - return false; - } - } - } - return true; - }); - // Update the map with the finalized result. - use_to_does_not_live_out[use] = does_not_live_out; - return does_not_live_out; -} - -// Filters out buffer uses that cannot use the cross-program prefetch due to -// aliasing with program output. -std::vector FindCrossProgramPrefetchUses( - absl::Span buffer_uses, - const HloAliasAnalysis& alias_analysis, const AliasInfo* alias_info) { - std::vector uses; - if (buffer_uses.empty()) { - return uses; - } - - const HloInstruction* root_instruction = buffer_uses.front() - .instruction->GetModule() - ->entry_computation() - ->root_instruction(); - - absl::flat_hash_map use_to_does_not_live_out; - for (const HloUse& use : buffer_uses) { - if (UseDoesNotLiveOut(use, alias_analysis, alias_info, root_instruction, - use_to_does_not_live_out)) { - uses.push_back(use); - } - } - return uses; -} - -bool IsCrossProgramPrefetchCandidate(const HloValue& value, - const HloAliasAnalysis& alias_analysis, - const AliasInfo* alias_info, - const Options& options) { - // Filter out values that alias with the entry computation root. - const HloBuffer& buffer = alias_analysis.GetBufferContainingValue(value); - const HloInstruction* root = alias_analysis.dataflow_analysis() - .module() - .entry_computation() - ->root_instruction(); - for (const HloPosition& position : buffer.ComputePositions()) { - if (position.instruction == root) { - return false; - } - } - std::vector uses = - FindCrossProgramPrefetchUses(value.GetUses(), alias_analysis, alias_info); - return value.defining_instruction()->parent() == - value.defining_instruction()->GetModule()->entry_computation() && - value.defining_instruction()->opcode() == HloOpcode::kParameter && - value.index().size() <= 1 && value.shape().IsArray() && - !uses.empty() && options.size_fn(value) <= options.max_size_in_bytes && - absl::c_all_of(uses, [&](const HloUse& use) { - const HloInstruction* inst = - use.instruction->operand(use.operand_number); - - // Skip the LooksLikeAnActivation test since we're testing the - // parent GTE/parameter and its children below. - if (inst->opcode() == HloOpcode::kBitcast && - ((inst->operand(0)->opcode() == HloOpcode::kGetTupleElement && - inst->operand(0)->operand(0)->opcode() == - HloOpcode::kParameter) || - inst->operand(0)->opcode() == HloOpcode::kParameter)) { - return true; - } - - return (inst->opcode() == HloOpcode::kGetTupleElement || - inst->opcode() == HloOpcode::kParameter) && - !LooksLikeAnActivation( - inst, options.cross_program_prefetch_permissive_mode); - }); -} - -bool IsUserAnnotatedCrossProgramPrefetch(const HloValue& value, - const Options& options) { - const HloInstruction* defining_instruction = value.defining_instruction(); - if (defining_instruction->parent() != - defining_instruction->GetModule()->entry_computation() || - defining_instruction->opcode() != HloOpcode::kParameter) { - return false; - } - const ComputationLayout& entry_computation_layout = - defining_instruction->GetModule()->entry_computation_layout(); - if (defining_instruction->parameter_number() >= - entry_computation_layout.parameter_count()) { - return false; - } - const Shape& shape = - entry_computation_layout - .parameter_layout(defining_instruction->parameter_number()) - .shape(); - return shape.has_layout() && - shape.layout().memory_space() == options.alternate_memory_space; -} - -MsaBufferInterval CreateMsaBufferInterval(const HloBuffer& buffer, - const HloValue* value, - const HloLiveRange& hlo_live_range, - const Options& options) { - MsaBufferInterval interval; - interval.buffer = value; - interval.size = options.size_fn(*value); - interval.start = 0; - interval.end = hlo_live_range.schedule_end_time(); - interval.colocations = {++buffer.values().begin(), buffer.values().end()}; - interval.need_allocation = true; - return interval; -} - -struct CrossProgramPrefetches { - std::vector prefetches; - std::vector candidates; -}; - -CrossProgramPrefetches FindCrossProgramPrefetches( - const HloAliasAnalysis& alias_analysis, const AliasInfo* alias_info, - const HloLiveRange& hlo_live_range, const Options& options) { - CrossProgramPrefetches cross_program_prefetches; - for (const HloBuffer& buffer : alias_analysis.buffers()) { - CHECK_GE(buffer.values().size(), 1); - const HloValue* value = buffer.values().at(0); - MsaBufferInterval buffer_interval = - CreateMsaBufferInterval(buffer, value, hlo_live_range, options); - if (IsUserAnnotatedCrossProgramPrefetch(*value, options)) { - cross_program_prefetches.prefetches.push_back(buffer_interval); - } else if (IsCrossProgramPrefetchCandidate(*value, alias_analysis, - alias_info, options)) { - cross_program_prefetches.candidates.push_back(buffer_interval); - } else if (MemorySpaceAssignmentUtils:: - DoesCrossProgramPrefetchBufferMatchAnyFilter( - options.msa_sort_order_overrides, buffer_interval)) { - cross_program_prefetches.candidates.push_back(buffer_interval); - } - } - - for (auto& prefetch : cross_program_prefetches.prefetches) { - VLOG(3) << "User annotated cross-program prefetch: " - << prefetch.buffer->ToString(); - } - - DefaultCrossProgramPrefetchBufferIntervalComparator default_comparator( - hlo_live_range, options.msa_sort_order_overrides); - BufferIntervalComparator* comparator = - (options.default_cross_program_prefetch_heuristic && - options.buffer_interval_comparator - ? options.buffer_interval_comparator - : &default_comparator); - absl::c_sort(cross_program_prefetches.candidates, - comparator->GetComparisonFunctor()); - - VLOG(3) << "Cross-program prefetch candidates: " - << cross_program_prefetches.candidates.size() - << ". Sorting criteria: " << comparator->DescribeComparisonCriteria(); - for (auto& candidate : cross_program_prefetches.candidates) { - VLOG(3) << "Cross-program prefetch candidate. Sorting criteria: " - << comparator->CriteriaToString(candidate) - << ". Candidate: " << candidate.buffer->ToString(); - } - return cross_program_prefetches; -} - -// Returns the conditional instruction that is the caller of the computation of -// which this instruction is the root, or nullptr if there is no such -// instruction. -HloInstruction* GetConditionalForBranchRoot(HloInstruction* branch_root) { - HloComputation* computation = branch_root->parent(); - if (computation->root_instruction() != branch_root) { - return nullptr; - } - for (HloInstruction* caller : computation->caller_instructions()) { - if (caller->opcode() == HloOpcode::kConditional) { - return caller; - } - } - return nullptr; -} - -// Returns the while instruction that is the caller of the body computation of -// which this instruction is the root, or nullptr if there is no such -// instruction. -HloInstruction* GetWhileForBodyRoot(HloInstruction* body_root) { - HloComputation* computation = body_root->parent(); - if (computation->root_instruction() != body_root) { - return nullptr; - } - for (HloInstruction* caller : computation->caller_instructions()) { - if (caller->opcode() == HloOpcode::kWhile && - caller->while_body() == computation) { - return caller; - } - } - return nullptr; -} - -// Returns the async pipelined while loop associated with 'instruction' (if -// 'instruction' is the while loop itself, the body parameter, or the body root -// tuple), or nullptr if there is none. -const HloInstruction* GetAsyncPipelinedWhileLoop( - const HloInstruction* instruction) { - if (instruction == nullptr) { - return nullptr; - } - if (instruction->opcode() == HloOpcode::kParameter) { - for (const HloInstruction* caller : - instruction->parent()->caller_instructions(HloOpcode::kWhile)) { - if (caller->while_body() == instruction->parent() && - IsAsyncPipelinedWhileLoop(caller)) { - return caller; - } - } - } else if (instruction->opcode() == HloOpcode::kWhile) { - if (IsAsyncPipelinedWhileLoop(instruction)) { - return instruction; - } - } else { - const HloComputation* comp = instruction->parent(); - if (comp != nullptr && instruction == comp->root_instruction() && - instruction->opcode() == HloOpcode::kTuple) { - HloInstruction* while_caller = - GetWhileForBodyRoot(const_cast(instruction)); - if (while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller)) { - return while_caller; - } - } - } - return nullptr; -} - -// Returns true if the position corresponds to loop-carried state (parameter or -// body root tuple) of an async pipelined while loop, or while/DUS/DS in an -// async pipelined while loop. -bool IsAsyncPipelinedWhilePosition(const HloPosition& pos) { - if (pos.instruction->opcode() == HloOpcode::kParameter) { - for (const HloInstruction* caller : - pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { - if (IsAsyncPipelinedWhileLoop(caller)) { - return true; - } - } - return false; - } - if (pos.instruction->opcode() == HloOpcode::kWhile) { - return IsAsyncPipelinedWhileLoop(pos.instruction); - } - if (pos.instruction->opcode() == HloOpcode::kDynamicUpdateSlice || - pos.instruction->opcode() == HloOpcode::kDynamicSlice) { - for (const HloInstruction* caller : - pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) { - if (IsAsyncPipelinedWhileLoop(caller)) { - return true; - } - } - return false; - } - HloComputation* comp = pos.instruction->parent(); - if (comp != nullptr && pos.instruction == comp->root_instruction() && - pos.instruction->opcode() == HloOpcode::kTuple) { - HloInstruction* while_caller = GetWhileForBodyRoot(pos.instruction); - return while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller); - } - return false; -} - -// Returns true if the computation is the body or condition of a while loop. -bool IsWhileBodyOrConditionComputation(const HloComputation* computation) { - if (computation == nullptr) { - return false; - } - for (const HloInstruction* caller : computation->caller_instructions()) { - if (caller->opcode() == HloOpcode::kWhile && - (caller->while_body() == computation || - caller->while_condition() == computation)) { - return true; - } - } - return false; -} - -int64_t GetShapeSizeBytes(const CostAnalysis* cost_analysis, - const Shape& shape) { - return cost_analysis ? cost_analysis->GetShapeSizeBytes(shape) : 0; -} - -void FindAliases(std::vector* allocation_values, - bool has_async_pipelined_while_loops = false) { - absl::flat_hash_map> - values_by_defining_inst; - for (AllocationValue& value : *allocation_values) { - values_by_defining_inst[value.position().instruction].push_back(&value); - } - auto maybe_add_alias_with_instruction = - [&](const HloInstruction* instruction, AllocationValue::Use* use, - std::optional expected_index = std::nullopt) { - auto aliased_values_it = values_by_defining_inst.find(instruction); - if (aliased_values_it != values_by_defining_inst.end()) { - for (const AllocationValue* aliased_value : - aliased_values_it->second) { - // Ensure that only matching tuple shape indexes are linked together - // as aliases when an expected shape index is specified (e.g. across - // while loop boundaries or async chain tuples). - if (expected_index.has_value() && - !aliased_value->position().index.empty() && - aliased_value->position().index != *expected_index) { - continue; - } - if (absl::c_find(use->aliases, aliased_value->position()) == - use->aliases.end()) { - VLOG(3) << "Adding aliasing for use " << use->hlo_use.ToString() - << " to " << aliased_value->ToShortString(); - use->aliases.push_back(aliased_value->position()); - } - } - } - }; - - for (AllocationValue& value : *allocation_values) { - for (AllocationValue::Use& use : value.uses()) { - // Find any aliases with the instruction itself (operand and output must - // alias), and any aliases with the parameters of called computations. - if (use.hlo_use.instruction->IsAsynchronous()) { - HloComputation* wrapped_computation = - use.hlo_use.instruction->async_wrapped_computation(); - if (use.hlo_use.instruction->opcode() == HloOpcode::kAsyncStart) { - // Bound operand i of async-start is placed at tuple index {0, i} of - // async-start's output shape. - ShapeIndex expected_async_index = {0, use.hlo_use.operand_number}; - expected_async_index.insert(expected_async_index.end(), - use.hlo_use.operand_index.begin(), - use.hlo_use.operand_index.end()); - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use, - expected_async_index); - if (use.hlo_use.instruction->async_wrapped_opcode() == - HloOpcode::kDynamicUpdateSlice && - use.hlo_use.operand_number == 0) { - ShapeIndex expected_dus_index = {1}; - expected_dus_index.insert(expected_dus_index.end(), - use.hlo_use.operand_index.begin(), - use.hlo_use.operand_index.end()); - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use, - expected_dus_index); - } - // Operands bound to async-start map directly to the initial - // parameters of the wrapped computation 1-to-1 by operand index. - if (use.hlo_use.operand_number < - wrapped_computation->num_parameters()) { - maybe_add_alias_with_instruction( - wrapped_computation->parameter_instruction( - use.hlo_use.operand_number), - &use); - } - } else if (use.hlo_use.instruction->opcode() == - HloOpcode::kAsyncUpdate) { - // For async-update, operand 0 is the async context tuple (the result - // of async-start or a prior async-update), which forwards matching - // tuple indices and does not correspond to a computation parameter. - // Operands 1, 2, ... are late-bound arguments mapping to wrapped - // computation parameters offset by the number of operands already - // bound by previous instructions in the async chain, and placed at - // tuple index {0, param_idx} of async-update's output shape. - if (use.hlo_use.operand_number == 0) { - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use, - use.hlo_use.operand_index); - } else if (use.hlo_use.operand_number >= 1) { - const HloInstruction* operand0 = - use.hlo_use.instruction->operand(0); - CHECK(operand0 != nullptr) << "Operand 0 is null for async update: " - << use.hlo_use.instruction->ToString(); - CHECK(operand0->opcode() == HloOpcode::kAsyncStart || - operand0->opcode() == HloOpcode::kAsyncUpdate) - << "Unexpected operand for async update: " - << use.hlo_use.instruction->ToString(); - std::vector previously_bound_operands = - hlo_instruction_utils::async::GetAsyncBoundOperands( - Cast(operand0)); - int previously_bound_operand_count = - previously_bound_operands.size(); - int param_idx = - use.hlo_use.operand_number - 1 + previously_bound_operand_count; - ShapeIndex expected_async_index = {0, param_idx}; - expected_async_index.insert(expected_async_index.end(), - use.hlo_use.operand_index.begin(), - use.hlo_use.operand_index.end()); - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use, - expected_async_index); - if (param_idx < wrapped_computation->num_parameters()) { - maybe_add_alias_with_instruction( - wrapped_computation->parameter_instruction(param_idx), &use); - } - } - } else { - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use); - } - } else if (use.hlo_use.instruction->opcode() == HloOpcode::kConditional) { - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use); - // A conditional instruction has branch_count + 1 operands: operand 0 is - // the branch index/predicate (which does not correspond to any branch - // computation parameter), and each operand i >= 1 is the single - // argument forwarded to branch_computation(i - 1)'s parameter 0. - if (use.hlo_use.operand_number >= 1) { - const HloComputation* branch_computation = - use.hlo_use.instruction->branch_computation( - use.hlo_use.operand_number - 1); - if (branch_computation->num_parameters() > 0) { - maybe_add_alias_with_instruction( - branch_computation->parameter_instruction(0), &use); - } - } - } else { - std::optional expected_index = std::nullopt; - if (use.hlo_use.instruction->opcode() == HloOpcode::kWhile) { - expected_index = use.hlo_use.operand_index; - } - maybe_add_alias_with_instruction(use.hlo_use.instruction, &use, - expected_index); - // For synchronous instructions with called computations (such as kCall - // and kWhile), caller operand i maps 1-to-1 to parameter i of each - // called computation. - for (const HloComputation* called_computation : - use.hlo_use.instruction->called_computations()) { - if (use.hlo_use.operand_number < - called_computation->num_parameters()) { - maybe_add_alias_with_instruction( - called_computation->parameter_instruction( - use.hlo_use.operand_number), - &use); - } - } - } - - // Special case for kWhile: the root of the body computation must alias as - // well. - if (use.hlo_use.instruction->opcode() == HloOpcode::kWhile) { - HloPosition root_alias{ - use.hlo_use.instruction->while_body()->root_instruction(), - use.hlo_use.operand_index}; - VLOG(3) << "Adding while body root aliasing for use " - << use.hlo_use.ToString() << " to " << root_alias; - use.aliases.push_back(root_alias); - } - - // Special case for conditionals - the output of a conditional op must - // alias with the branch computation outputs. - HloInstruction* conditional_instruction = - GetConditionalForBranchRoot(use.hlo_use.instruction); - if (conditional_instruction != nullptr && - use.hlo_use.instruction->opcode() == HloOpcode::kTuple) { - // We only need to add a use alias if the branch root is a tuple, - // because a tuple is a use and any other instruction would be a - // definition or a position. - ShapeIndex index = use.hlo_use.operand_index; - index.push_front(use.hlo_use.operand_number); - HloPosition conditional_output_position{conditional_instruction, index}; - VLOG(1) << "Add use alias for counditional output position " - << conditional_output_position.ToString() << " to use " - << use.hlo_use.ToString(); - use.aliases.push_back(conditional_output_position); - } - - // Special case for async pipelined while loops: an instruction - // producing a loop return operand that feeds into the body root tuple - // must alias with both the while body parameter(0) and the caller while - // instruction at the corresponding tuple index. - if (has_async_pipelined_while_loops) { - HloInstruction* while_instruction = - GetWhileForBodyRoot(use.hlo_use.instruction); - if (while_instruction != nullptr && - IsAsyncPipelinedWhileLoop(while_instruction) && - use.hlo_use.instruction->opcode() == HloOpcode::kTuple && - while_instruction->while_body()->num_parameters() > 0) { - ShapeIndex index = use.hlo_use.operand_index; - index.push_front(use.hlo_use.operand_number); - HloInstruction* parameter_instruction = - while_instruction->while_body()->parameter_instruction(0); - HloPosition parameter_position{parameter_instruction, index}; - HloPosition while_output_position{while_instruction, index}; - VLOG(1) - << "Add use aliases for while body root position to parameter " - << parameter_position.ToString() << " and while output " - << while_output_position.ToString() << " for use " - << use.hlo_use.ToString(); - use.aliases.push_back(parameter_position); - use.aliases.push_back(while_output_position); - } - } - } - } -} - -} // namespace - -// Helper to inspect the async wrapped opcode of a pipelined while loop -// position. -std::optional GetAsyncPipelinedWhileWrappedOpcode( - const HloInstruction* while_instr, const HloPosition& pos, - const HloAliasAnalysis& alias_analysis) { - if (while_instr == nullptr || while_instr->opcode() != HloOpcode::kWhile || - pos.index.empty()) { - return std::nullopt; - } - - int64_t tuple_idx = pos.index[0]; - ShapeIndex base_index = {tuple_idx, 0, 0}; - if (!ShapeUtil::IndexIsValid(pos.instruction->shape(), base_index)) { - return std::nullopt; - } - - std::vector buffers = - alias_analysis.ComputeBuffersAt(pos.instruction, base_index); - const HloInstruction* async_start = nullptr; - for (const HloBuffer* buffer : buffers) { - for (const HloPosition& p : buffer->ComputePositions()) { - if (p.instruction->opcode() == HloOpcode::kAsyncStart && - p.instruction->parent() == while_instr->while_body()) { - async_start = p.instruction; - } - } - } - - if (async_start == nullptr) { - return std::nullopt; - } - return async_start->async_wrapped_opcode(); -} - -std::optional GetAsyncPipelinedWhileWrappedOpcode( - const HloInstruction* while_instr, int64_t tuple_idx, - const HloAliasAnalysis& alias_analysis) { - if (while_instr == nullptr || while_instr->opcode() != HloOpcode::kWhile) { - return std::nullopt; - } - return GetAsyncPipelinedWhileWrappedOpcode( - while_instr, - HloPosition{const_cast(while_instr), {tuple_idx}}, - alias_analysis); -} - -// Returns true if the position in an async pipelined while loop corresponds to -// a buffer that is intended to reside in alternate memory (e.g., prefetched -// dynamic-slice output, or dynamic-update-slice update slice). Base tensors of -// dynamic-slice or dynamic-update-slice and dynamic-update-slice outputs -// reside in default memory (HBM) on TPU and return false. -bool IsAsyncPipelinedWhileAlternateMemoryPosition( - const HloPosition& pos, const HloAliasAnalysis& alias_analysis) { - if (pos.instruction->IsAsynchronous()) { - HloOpcode wrapped_opcode = pos.instruction->async_wrapped_opcode(); - if (wrapped_opcode == HloOpcode::kDynamicSlice) { - if (pos.instruction->opcode() == HloOpcode::kAsyncDone) { - return pos.index.empty(); - } - if (pos.instruction->opcode() == HloOpcode::kAsyncStart) { - return pos.index == ShapeIndex({1}); - } - } - // DynamicUpdateSlice base tensor and output reside in default memory. - return false; - } - - if (pos.index.empty()) { - return false; - } - - const HloInstruction* while_instr = - GetAsyncPipelinedWhileLoop(pos.instruction); - if (while_instr == nullptr) { - return false; - } - - int64_t tuple_idx = pos.index[0]; - ShapeIndex base_index = {tuple_idx, 0, 0}; - if (!ShapeUtil::IndexIsValid(pos.instruction->shape(), base_index)) { - return false; - } - - std::optional wrapped_opcode = - GetAsyncPipelinedWhileWrappedOpcode(while_instr, pos, alias_analysis); - if (!wrapped_opcode.has_value()) { - return false; - } - - if (*wrapped_opcode == HloOpcode::kDynamicUpdateSlice) { - // In an async dynamic-update-slice bundle: - // {tuple_idx, 0, 0} is the base tensor -> default memory - // {tuple_idx, 0, 1} is the update slice -> alternate memory - // {tuple_idx, 1} is the dynamic-update-slice output (aliased to base) -> - // default memory - return pos.index == ShapeIndex({tuple_idx, 0, 1}); - } - - if (*wrapped_opcode == HloOpcode::kDynamicSlice) { - // In an async dynamic-slice bundle: - // {tuple_idx, 0, 0} is the base tensor -> default memory - // {tuple_idx, 1} or {tuple_idx} is the slice output -> alternate memory - return pos.index == ShapeIndex({tuple_idx, 1}) || - pos.index == ShapeIndex({tuple_idx}); - } - - return false; -} - -bool TimeInterval::operator<(const TimeInterval& other) const { - return std::forward_as_tuple(inclusive_start_time, inclusive_end_time) < - std::forward_as_tuple(other.inclusive_start_time, - other.inclusive_end_time); -} - -std::string TimeInterval::ToString() const { - return absl::StrCat("[", inclusive_start_time, ", ", inclusive_end_time, "]"); -} - -bool MsaAlgorithm::IsIntervalPinnedToAlternateMemory( - const MsaBufferInterval& interval) const { - const Shape& shape = interval.buffer->shape(); - return shape.has_layout() && - shape.layout().memory_space() == options_.alternate_memory_space; -} - -bool MsaAlgorithm::MatchesPrefetchContext( - const PrefetchContext& context, absl::string_view producer_name, - ShapeIndex producer_shape_index, absl::string_view consumer_name) const { - return context.request->use->hlo_use.instruction->name() == consumer_name && - context.request->allocation_value->position().instruction->name() == - producer_name && - context.request->allocation_value->position().index == - producer_shape_index; -} - -MsaAlgorithm::MsaAlgorithm(HloModule* module, AllocationSequence* allocations, - const Options& options, - const HloAliasAnalysis& alias_analysis, - const AliasInfo* alias_info, - const HloLiveRange& hlo_live_range) - : GlobalDecreasingSizeBestFitHeap( - options.alignment_in_bytes, - /*type=*/kSpatial, /*buffer_interval_compare=*/nullptr, - (options.sliced_prefetch_options.max_slices() > - options.sliced_prefetch_options - .all_slice_time_permutations_threshold() - ? SliceTimePermutationIterator::Ty::kPreferred - : SliceTimePermutationIterator::Ty::kAll)), - module_(module), - allocations_(allocations), - options_(options), - alias_analysis_(alias_analysis), - alias_info_(alias_info), - hlo_live_range_(hlo_live_range), - peak_memory_usage_(hlo_live_range.schedule_end_time() + 1) { - // Override buffer interval compare if provided. - auto comparison_function = GetSpatialBufferIntervalCompare(); - if (options.buffer_interval_comparator) { - comparison_function = - options.buffer_interval_comparator->GetComparisonFunctor(); - } - - // Prioritize pinned buffers in the buffer interval order. - buffer_interval_compare_ = - [this, comparison_function = std::move(comparison_function)]( - const MsaBufferInterval& a, const MsaBufferInterval& b) { - bool is_a_pinned = IsIntervalPinnedToAlternateMemory(a); - bool is_b_pinned = IsIntervalPinnedToAlternateMemory(b); - if (is_a_pinned && !is_b_pinned) { - return true; - } - if (!is_a_pinned && is_b_pinned) { - return false; - } - return comparison_function(a, b); - }; - - call_graph_ = CallGraph::Build(&alias_analysis_.dataflow_analysis().module()); - - absl::flat_hash_map - instruction_shape_hash_cache; - std::vector initial_resources(hlo_live_range.schedule_end_time(), 1.0); - if (options.cost_analysis) { - const std::vector& flattened_instructions = - hlo_live_range.flattened_instruction_sequence().instructions(); - int64_t farthest_bandwidth_limiting_async_done_found = -1; - float current_bw_adjustment_factor = 1.0f; - for (int i = 0; i < flattened_instructions.size(); ++i) { - const HloInstruction* inst = flattened_instructions[i]; - std::optional bw_adjustment_factor = - options.async_instruction_bw_adjustment_factor_fn(inst); - if (bw_adjustment_factor.has_value()) { - CHECK_EQ(inst->users().size(), 1); - current_bw_adjustment_factor = bw_adjustment_factor.value(); - farthest_bandwidth_limiting_async_done_found = std::max( - farthest_bandwidth_limiting_async_done_found, - hlo_live_range.instruction_schedule().at(inst->users()[0])); - } - if (!MemorySpaceAssignmentUtils::IsInstructionOnConfiguredExecThread( - *inst, options.execution_threads) || - inst->opcode() == HloOpcode::kWhile || - inst->opcode() == HloOpcode::kConditional) { - initial_resources[i] = 0; - } else { - initial_resources[i] = - options.cost_analysis->GetInstructionElapsed(*inst); - if (options_.use_repeated_instance_for_preferred_prefetch_time || - options_.memory_bound_loop_optimizer_options.enabled()) { - uint64_t fingerprint = - absl::HashOf(MsaInstructionFingerprint( - inst, &instruction_shape_hash_cache)); - fingerprint_map_[inst] = fingerprint; - repeated_inst_map_[fingerprint].push_back(inst); - } - } - if (i <= farthest_bandwidth_limiting_async_done_found) { - VLOG(2) << "Adjusting initial_resource[" << i << "] " - << initial_resources[i] << " " << inst->name() - << " because of overlap with bandwidth limiting async start " - "and done instructions."; - initial_resources[i] *= current_bw_adjustment_factor; - } - VLOG(2) << "Initial resource[" << i << "] = " << initial_resources[i] - << " (" << inst->name() << ")"; - } - } - prefetch_async_copy_resource_ = AsynchronousCopyResource(initial_resources); - eviction_async_copy_resource_ = AsynchronousCopyResource(initial_resources); -} - -namespace { - -// Finds the allocation value for operand 0 of async updates and dones. We make -// sure these uses are associated with the AllocationValue whose position -// is the direct source of the use. -AllocationValue* FindAllocationValueForAsyncOperationStateUse( - const HloUse& use, - absl::Span candidate_allocation_values) { - CHECK_EQ(use.operand_number, 0); - HloPosition source_position = GetNonTrivialSourcePosition( - HloPosition{use.instruction->mutable_operand(0), use.operand_index}); - for (AllocationValue& allocation_value : candidate_allocation_values) { - if (allocation_value.position() == source_position) { - return &allocation_value; - } - } - return nullptr; -} - -// Finds the latest valid AllocationValue defined before use_time in the same -// computation as the use. -// -// candidate_allocation_values contains all AllocationValue objects created for -// the non-trivial positions (excluding GTE, Tuple, and Bitcast) of the single -// HloValue currently being processed, sorted chronologically by definition -// time. -// -// For dataflow uses, we search backward to find the most recent definition -// before use_time in the use's computation. If skip_async_state_definitions is -// true, definitions from async start/update instructions are skipped because -// they represent in-flight bundled async state that is only finalized when the -// corresponding async-done executes. -// -// If we can't find a valid AllocationValue, returns nullptr. -AllocationValue* FindLatestAllocationValueForUse( - absl::Span candidate_allocation_values, int64_t use_time, - HloComputation* use_computation, - const absl::flat_hash_map& - instruction_schedule, - bool skip_async_state_definitions) { - for (auto it = candidate_allocation_values.rbegin(); - it != candidate_allocation_values.rend(); ++it) { - AllocationValue* allocation_value = &(*it); - const HloInstruction* defining_instruction = - allocation_value->position().instruction; - - // Skip definitions from different computations. - if (allocation_value->computation() != use_computation) { - continue; - } - - // Skip definitions that are after the use. - int64_t definition_time = instruction_schedule.at(defining_instruction); - if (definition_time >= use_time) { - continue; - } - - // Async start/update instructions produce in-flight bundled async state - // that is not finalized until the corresponding async-done executes. - // Regular uses must not associate directly with these unfinalized state - // definitions. - if (skip_async_state_definitions && - IsAsyncOperationStateDefinition(defining_instruction)) { - continue; - } - - return allocation_value; - } - return nullptr; -} - -} // namespace - -// Finds the matching AllocationValue for a given HloUse. -AllocationValue* MsaAlgorithm::FindAllocationValueForUse( - const HloUse& use, absl::Span candidate_allocation_values, - int64_t use_time) const { - if (IsAsyncOperationStateUse(use)) { - // Case A: Async operation state uses. Note, this case covers uses of - // bundled async state in traditional serial async chains and in pipelined - // while loops. - return FindAllocationValueForAsyncOperationStateUse( - use, candidate_allocation_values); - } - - // Case B: All other uses. These uses use the latest allocation value for the - // operand, potentially skipping async state definitions. - const absl::flat_hash_map& - instruction_schedule = hlo_live_range_.instruction_schedule(); - - HloInstruction* while_for_body_root = GetWhileForBodyRoot(use.instruction); - - bool is_pipelined_while_async_start_use = - (use.instruction->opcode() == HloOpcode::kWhile && - IsAsyncPipelinedWhileLoop(use.instruction)) || - (while_for_body_root != nullptr && - IsAsyncPipelinedWhileLoop(while_for_body_root) && - use.instruction->opcode() == HloOpcode::kTuple); - - return FindLatestAllocationValueForUse( - candidate_allocation_values, use_time, use.instruction->parent(), - instruction_schedule, - /*skip_async_state_definitions=*/!is_pipelined_while_async_start_use); -} - -void MsaAlgorithm::CreateAllocationValues( - const MsaBufferInterval& buffer_interval, - std::vector& allocation_values) const { - const HloValue* value = buffer_interval.buffer; - VLOG(3) << "Creating AllocationValues"; - - // Find and sort all non-trivial (excluding GTE, Tuple, and bitcast) - // positions. We create an AllocationValue object for each non-trivial - // position. And for each AllocationValue object, we create an - // AllocationSequence consisting of one or more Allocation objects.The reason - // why we exclude the trivial positions from AllocationValue is because - // Allocation objects have special support for tuples and bitcasts. - const absl::flat_hash_map& - instruction_schedule = hlo_live_range_.instruction_schedule(); - std::vector positions; - for (const HloPosition& position : value->positions()) { - const HloInstruction* instruction = position.instruction; - if (instruction->opcode() != HloOpcode::kGetTupleElement && - instruction->opcode() != HloOpcode::kTuple && - instruction->opcode() != HloOpcode::kBitcast) { - positions.push_back(position); - } - } - absl::c_stable_sort(positions, - [&](const HloPosition& pos1, const HloPosition& pos2) { - return instruction_schedule.at(pos1.instruction) < - instruction_schedule.at(pos2.instruction); - }); - - // Create an AllocationValue for each non-trivial position. - int beginning_idx = allocation_values.size(); - for (int i = 0; i < positions.size(); ++i) { - const HloPosition& position = positions.at(i); - allocation_values.emplace_back(value, position, buffer_interval.size); - } - - std::vector uses(value->GetUses().begin(), value->GetUses().end()); - absl::c_stable_sort(uses, [&](const HloUse& use1, const HloUse& use2) { - return instruction_schedule.at(use1.instruction) < - instruction_schedule.at(use2.instruction); - }); - - // Associate each use with an AllocationValue. Each AllocationValue contains a - // position and uses in the same computation. Furthermore, if the original - // HloValue had multiple non-trivial positions in the same computation, those - // will get their own AllocationValue as well. We split these HloValues so - // that when we insert CopyStart/CopyDone in CopyAllocation::Process, they - // point to the latest position. We then replace the operand of the use with - // CopyStart/CopyDone with an operand of the latest position. - absl::Span candidate_allocation_values = - absl::MakeSpan(allocation_values).subspan(beginning_idx); - for (const HloUse& use : uses) { - int64_t use_time = instruction_schedule.at(use.instruction); - AllocationValue* allocation_value_for_use = - FindAllocationValueForUse(use, candidate_allocation_values, use_time); - - CHECK(allocation_value_for_use != nullptr) - << "Failed to find AllocationValue for use: " << use.ToString(); - allocation_value_for_use->AddUse(use, use_time); - } - - // This loop marks allocation values as contiguous, if they are part of async - // operations, and logs all created allocation values. - for (int i = beginning_idx; i < allocation_values.size(); ++i) { - AllocationValue& allocation_value = allocation_values.at(i); - if (IsAsyncOperationStateDefinition( - allocation_value.position().instruction)) { - for (const AllocationValue::Use& use : allocation_value.uses()) { - HloInstruction* use_instruction = use.hlo_use.instruction; - CHECK(use_instruction->opcode() == HloOpcode::kAsyncUpdate || - HloDataflowAnalysis::IsAsynchronousOperationDone( - use_instruction->opcode()) || - HloDataflowAnalysis::IsAsynchronousOperationStart( - use_instruction->opcode()) || - use_instruction->opcode() == HloOpcode::kTuple || - use_instruction->opcode() == HloOpcode::kGetTupleElement || - use_instruction->opcode() == HloOpcode::kWhile || - use_instruction->opcode() == HloOpcode::kConditional) - << "Unexpected use_instruction opcode: " - << HloOpcodeString(use_instruction->opcode()) << " (" - << use_instruction->ToString() << ") for async def: " - << allocation_value.position().instruction->ToString(); - } - } - bool is_async_operation_state = - IsAsyncOperationStateDefinition( - allocation_value.position().instruction) || - (has_async_pipelined_while_loops_ && - IsBufferAliasedToAsyncPipelinedWhileLoop(allocation_value.value()) && - (IsAsyncPipelinedWhilePosition(allocation_value.position()) || - absl::c_any_of(allocation_value.uses(), - [](const AllocationValue::Use& use) { - return IsAsyncOperationStateUse(use.hlo_use); - }))); - // Requiring contiguous allocation ensures that buffers representing inputs - // and outputs to the async computation maintain temporal contiguity - // (preventing MSA from evicting them to default memory). - // - // For async pipelined while loops where the async start/update is in the - // while loop body and the done is in the next iteration, the uses - // can include kTuple/kGetTupleElement passing the async state out of the - // while loop body, or kWhile/kConditional when unrolling or lowering. - // - // However, this does not guarantee temporal contiguousness for temporary - // buffers created *inside* the async computation. Such buffers must be - // allocated outside the async operation and passed in as `kParameter`s to - // ensure they are handled correctly. - if (is_async_operation_state) { - VLOG(3) << "Mark " << allocation_value.ToShortString() - << " to require contiguous allocation because it is an async " - "operation state (async start/update definition or async " - "done/update use)."; - allocation_value.set_requires_contiguous_allocation(true); - } else if (options_.position_requires_contiguous_allocation_fn( - allocation_value.position())) { - VLOG(3) << "Mark " << allocation_value.ToShortString() - << " to require contiguous allocation because of options."; - allocation_value.set_requires_contiguous_allocation(true); - } - VLOG(3) << "Created allocation value: " - << allocation_values.at(i).ToString(); - } -} - -void MsaAlgorithm::FindAliasesForTesting( - std::vector* allocation_values, - bool has_async_pipelined_while_loops) { - FindAliases(allocation_values, has_async_pipelined_while_loops); -} - -void MsaAlgorithm::ExtendScopedAlternateMemoryAllocations() { - VLOG(1) << "Starting vmem expansion"; - - // Iterate through all scoped allocations and try to expand them to the - // largest contiguous open space available. - for (std::unique_ptr& allocation : *allocations_) { - if (!allocation->is_scoped_allocation()) { - continue; - } - - // Find the set of nodes that are live during allocation. - std::vector live_nodes = interval_tree_.ChunksOverlappingInTime( - allocation->start_time(), allocation->end_time()); - absl::c_sort(live_nodes, [](const Chunk& a, const Chunk& b) { - return a.offset < b.offset; - }); - - // Loop over live_nodes to compute 2 things: - // 1. The largest contiguous free chunk (biggest_free_chunk) - // 2. The largest chunk we can get by moving the start time of the scoped - // allocation earlier (i.e., to max_end_before_scoped_allocation), and - // the end time later (i.e., to min_offset_after_scoped_allocation). - int64_t min_offset_after_scoped_allocation = available_heap_size(); - int64_t max_end_before_scoped_allocation = 0; - Chunk biggest_free_chunk = Chunk::FromOffsetSize(0, 0); - for (int i = 0; i < live_nodes.size(); ++i) { - const Chunk& chunk = live_nodes[i]; - if (allocation->chunk().chunk_end() <= chunk.offset) { - min_offset_after_scoped_allocation = - std::min(min_offset_after_scoped_allocation, chunk.offset); - } - if (allocation->chunk().offset >= chunk.chunk_end()) { - max_end_before_scoped_allocation = - std::max(max_end_before_scoped_allocation, chunk.chunk_end()); - } - - Chunk next_free_chunk = Chunk::FromOffsetEnd( - GetAlignedOffset(chunk.chunk_end(), options_.alignment_in_bytes), - (i + 1) < live_nodes.size() ? live_nodes[i + 1].offset - : available_heap_size()); - if (next_free_chunk.size > biggest_free_chunk.size) { - biggest_free_chunk = next_free_chunk; - } - } - - Chunk proposed_extended_chunk = - Chunk::FromOffsetEnd(GetAlignedOffset(max_end_before_scoped_allocation, - options_.alignment_in_bytes), - min_offset_after_scoped_allocation); - - // Check if we should extend the boundaries of the scoped allocation or - // move it. - Chunk proposed_chunk = allocation->chunk(); - std::string source; - if (proposed_extended_chunk.size > proposed_chunk.size) { - proposed_chunk = proposed_extended_chunk; - source = "extended"; - } - if (biggest_free_chunk.size > proposed_chunk.size) { - proposed_chunk = biggest_free_chunk; - source = "free"; - } - if (source.empty()) { - VLOG(3) << "Could not move the scoped allocation for " - << allocation->defining_position().ToString() - << "; Current fragmentation: " << - [&]() { - int64_t occupied_size = 0; - for (const Chunk& chunk : live_nodes) { - occupied_size += chunk.size; - } - double fragmentation = - static_cast(available_heap_size() - occupied_size) / - static_cast(available_heap_size()); - return 100.0 * fragmentation; - }() << "%"; - continue; - } - - VLOG(1) << "Moving the scoped allocation for " - << allocation->defining_position().ToString() << " from " - << allocation->chunk().ToString() << " to " - << proposed_chunk.ToString() << " (" << source - << "); Size increase: " - << (100.0 * - static_cast(proposed_chunk.size - - allocation->chunk().size) / - static_cast(allocation->chunk().size)) - << "%"; - - // Update the allocation. We don't need to update result_.chunk_map. It's - // not used by MSA. - *allocation->mutable_chunk() = proposed_chunk; - result_.UpdatedHeapSize(proposed_chunk); - } -} - -std::string MsaAlgorithm::RequiredMemoryAssignment::SourceToString( - Source source) { - switch (source) { - case Source::kAliasedUse: - return "Aliased use"; - case Source::kUseNotAllowedInAlternateMemory: - return "Use not allowed in alternate memory."; - case Source::kInefficientSite: - return "Inefficient site"; - case Source::kLoopOptimizedParameterInWhileLoop: - return "Loop optimized parameter in while loop"; - case Source::kConditionalPhiOutput: - return "Conditional phi output"; - case Source::kConditionalComputationOutput: - return "Conditional computation output"; - case Source::kPositionNotAllowedInAlternateMemory: - return "Position not allowed in alternate memory."; - case Source::kProgramOutput: - return "Program output."; - case Source::kConstantInstruction: - return "Constant instruction."; - case Source::kProgramInput: - return "Program input."; - case Source::kBlockPrefetchSourceBuffer: - return "Block prefetch source buffer."; - case Source::kBlockPrefetchSourceBufferUse: - return "Block prefetch source buffer use."; - } -} - -std::string MsaAlgorithm::RequiredMemoryAssignment::ToString() const { - std::string memory_space_str = - memory_space == MemorySpace::kDefault ? "def" : "alt"; - std::string offset_str = - offset == nullptr ? "null" : absl::StrCat(offset->offset); - - return absl::StrCat( - "RequiredMemoryAssignment(memory_space=", memory_space_str, - ", time=", time, ", offset=", offset_str, - ", message=", SourceToString(required_assignment_source), ")"); -} - -std::vector MsaAlgorithm::GetSortedColocatedIntervals( - const MsaBufferInterval& interval) const { - std::vector colocated_intervals; - std::vector worklist = {&interval}; - while (!worklist.empty()) { - const MsaBufferInterval* item = worklist.back(); - worklist.pop_back(); - colocated_intervals.push_back(item); - for (const HloValue* buffer_colocated : item->colocations) { - worklist.push_back(&buffer_intervals_.at(buffer_colocated)); - } - } - - absl::c_stable_sort(colocated_intervals, [&](const MsaBufferInterval* x, - const MsaBufferInterval* y) { - return std::make_pair(x->start, x->end) < std::make_pair(y->start, y->end); - }); - return colocated_intervals; -} - -bool MsaAlgorithm::IsUseAllowedInAlternateMemory(const HloUse& use) const { - if (!options_.is_use_allowed_in_alternate_mem_fn(use)) { - return false; - } - return true; -} - -bool MsaAlgorithm::IsWhileLoopUseRequiredInDefaultMemory( - const HloUse& use) const { - // Check if there is a required assignment for the while loop output. - if (use.instruction->opcode() == HloOpcode::kWhile) { - HloValue* while_value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - use.instruction, use.operand_index); - int64_t while_time = - hlo_live_range_.instruction_schedule().at(use.instruction); - auto existing_required_assignment = - RequiredMemoryAssignmentAt(while_value, while_time); - if (existing_required_assignment && - existing_required_assignment->memory_space == MemorySpace::kDefault) { - VLOG(4) << "While allocation not allowed in alternate memory because " - "there is a required default memory assignment."; - return true; - } - } - return false; -} - -bool MsaAlgorithm::IsWhileLoopUseBeneficialInAlternateMemory( - const HloUse& use) const { - if (use.instruction->opcode() != HloOpcode::kWhile) { - return true; - } - HloComputation* while_body = use.instruction->while_body(); - - // We don't want to allocate this buffer in alternate memory if it will be - // evicted anyway. Find out if it has an early use or a late definition that - // would make sense to keep it in the alternate memory. - HloValue* parameter_value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - while_body->parameter_instruction(0), use.operand_index); - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - int64_t parameter_time = - instruction_schedule.at(while_body->parameter_instruction(0)); - int64_t root_time = instruction_schedule.at(while_body->root_instruction()); - int64_t min_use_time = root_time; - for (const HloUse& parameter_use : parameter_value->GetUses()) { - int64_t use_time = instruction_schedule.at(parameter_use.instruction); - if (parameter_use.instruction->opcode() != HloOpcode::kGetTupleElement && - parameter_use.instruction->opcode() != HloOpcode::kTuple && - parameter_use.instruction->opcode() != HloOpcode::kBitcast && - use_time > parameter_time) { - min_use_time = std::min(min_use_time, use_time); - } - } - // If there is no use of this buffer inside the while loop, there is no need - // to allocate it in the loop. - if (min_use_time == root_time) { - VLOG(4) << "While allocation not allowed in alternate memory. " - << "use time = " << min_use_time << ", root time = " << root_time; - return false; - } - const Shape& shape = parameter_value->shape(); - // Allow the buffer in alternate memory if the buffer has a short live range - // either at the beginning or end of the while loop body. - int64_t shape_size = GetShapeSizeBytes(options_.cost_analysis, shape); - if (!options_.prefetch_interval_picker->CanAllocateInAlternateMemoryNoCopy( - shape_size, parameter_time, min_use_time)) { - VLOG(4) << "While allocation not allowed in alternate memory. " - << "use time = " << min_use_time << ", root time = " << root_time; - return false; - } - return true; -} - -bool MsaAlgorithm::IsConditionalUseBeneficialInAlternateMemory( - const AllocationValue& value, const HloUse& use) const { - if (use.instruction->opcode() != HloOpcode::kConditional) { - return true; - } - // For any use of this conditional (the same value might be passed into - // multiple called computations), determine if the parameter->first use - // dependency is short. - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - int64_t conditional_time = instruction_schedule.at(use.instruction); - for (const AllocationValue::Use& other_use : value.uses()) { - if (other_use.hlo_use.instruction != use.instruction) { - continue; - } - // Operand 0 is not passed into the computation. - if (other_use.hlo_use.operand_number == 0) { - continue; - } - HloComputation* called_computation = - use.instruction->called_computations().at( - other_use.hlo_use.operand_number - 1); - const HloInstruction* parameter_instruction = - called_computation->parameter_instruction(0); - HloValue* parameter_value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - parameter_instruction, other_use.hlo_use.operand_index); - int64_t parameter_time = instruction_schedule.at(parameter_instruction); - int64_t min_use_time = conditional_time; - for (const HloUse& parameter_use : parameter_value->GetUses()) { - if (parameter_use.instruction->parent() == called_computation && - parameter_use.instruction->opcode() != HloOpcode::kGetTupleElement && - parameter_use.instruction->opcode() != HloOpcode::kTuple && - parameter_use.instruction->opcode() != HloOpcode::kBitcast) { - min_use_time = std::min( - min_use_time, instruction_schedule.at(parameter_use.instruction)); - } - } - int64_t shape_size = - GetShapeSizeBytes(options_.cost_analysis, parameter_value->shape()); - if (options_.prefetch_interval_picker->CanAllocateInAlternateMemoryNoCopy( - shape_size, parameter_time, min_use_time)) { - VLOG(4) << "Conditional allocation allowed in alternate memory for " - "computation = " - << called_computation->name() - << ", parameter time = " << parameter_time - << ", min use time = " << min_use_time; - return true; - } - VLOG(4) << "Conditional allocation not allowed in alternate memory for " - "computation = " - << called_computation->name() - << ", parameter time = " << parameter_time - << ", min use time = " << min_use_time; - } - return false; -} - -namespace { -// Columns in buffer information: -// buffer_id: int. This value can be used to match the allocation in -// allocation information. -// buffer_name: string. -// alt_mem_benefit: float. Roughly corresponds to how much the cost analysis -// thought it would be beneficial to put this in the alternate memory. The -// higher the value, the more it is memory bound. -// size: int. In bytes. -// definition_time: int. Logical time this value was defined in the schedule. -// use_times: string. This is a semicolon-separated list of integers for all -// the use times. -// use_names: string. This is a semicolon-separated list of string -// representation of uses. -// is_scoped: int. A value of 1 indicates that the buffer is a scoped -// allocation. -constexpr absl::string_view kBufferInfoColumnNames = - "buffer_id,buffer_name,alt_mem_benefit,size,definition_time,use_times,use_" - "names,is_scoped"; -} // namespace - -void MsaAlgorithm::AppendBufferInfoDebugString( - const MsaBufferInterval& interval, std::string* debug_str) const { - if (debug_str->empty()) { - // Append the column names. - absl::StrAppend(debug_str, kBufferInfoColumnNames, "\n"); - } - const HloBuffer& buffer = - alias_analysis_.GetBufferContainingValue(*interval.buffer); - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - int64_t definition_time = - instruction_schedule.at(interval.buffer->defining_position().instruction); - std::vector> uses; - for (const HloValue* value : buffer.values()) { - for (const HloUse& use : value->GetUses()) { - uses.push_back( - {instruction_schedule.at(use.instruction), use.ToString()}); - } - } - absl::c_sort(uses); - std::vector use_times; - std::vector use_names; - use_times.reserve(uses.size()); - use_names.reserve(uses.size()); - for (const auto& use : uses) { - use_times.push_back(use.first); - use_names.push_back(use.second); - } - - absl::StrAppend(debug_str, buffer.id(), ","); - absl::StrAppend(debug_str, "\"", interval.buffer->ToShortString(), "\","); - auto alternate_memory_benefit = - options_.prefetch_interval_picker->BufferIntervalAlternateMemoryBenefit( - interval); - absl::StrAppend( - debug_str, alternate_memory_benefit ? *alternate_memory_benefit : 0, ","); - absl::StrAppend(debug_str, interval.size, ","); - absl::StrAppend(debug_str, definition_time, ","); - absl::StrAppend(debug_str, "\"", absl::StrJoin(use_times, ";"), "\","); - absl::StrAppend(debug_str, "\"", absl::StrJoin(use_names, ";"), "\","); - absl::StrAppend(debug_str, "0"); // is_scoped - absl::StrAppend(debug_str, "\n"); -} - -void MsaAlgorithm::AppendScopedAllocationBufferInfoDebugString( - const HloInstruction* instruction, int64_t time, int64_t size, - std::string& debug_str) const { - if (debug_str.empty()) { - // Append the column names. - absl::StrAppend(&debug_str, kBufferInfoColumnNames, "\n"); - } - int64_t buffer_id = alias_analysis_.GetUniqueBufferAt(instruction).id(); - - // As a convention, we use negative values for scoped allocations. - absl::StrAppend(&debug_str, -buffer_id, ","); - absl::StrAppend(&debug_str, "\"scoped allocation for ", instruction->name(), - "\","); - absl::StrAppend(&debug_str, 0, ","); // alt_mem_benefit - absl::StrAppend(&debug_str, size, ","); - absl::StrAppend(&debug_str, time, ","); - absl::StrAppend(&debug_str, "\"\","); // use_times - absl::StrAppend(&debug_str, "\"\","); // use_names - absl::StrAppend(&debug_str, "1"); // is_scoped - absl::StrAppend(&debug_str, "\n"); -} - -void MsaAlgorithm::AppendAllocationInfoDebugString( - const Allocation& allocation, std::string& debug_str) const { - // Columns in allocation information: - // buffer_id: int. This value can be used the match with buffer info. - // size: int. In bytes. - // offset: int. In bytes. - // start_time: int. Logical start time of the allocation. - // end_time: int. Logical end time of the allocation. - if (debug_str.empty()) { - // Append the column names. - absl::StrAppend(&debug_str, "buffer_id,size,offset,start_time,end_time\n"); - } - if (allocation.memory_space() == MemorySpace::kAlternate) { - const HloPosition& position = allocation.defining_position(); - int64_t buffer_id = position.instruction->unique_id(); - if (alias_analysis_.ComputeBuffersAt(position.instruction, position.index) - .size() == 1) { - buffer_id = alias_analysis_ - .GetUniqueBufferAt(position.instruction, position.index) - .id(); - } - // As a convention, we use negative values for scoped allocations. - absl::StrAppend(&debug_str, - allocation.is_scoped_allocation() ? -buffer_id : buffer_id, - ","); - absl::StrAppend(&debug_str, allocation.chunk().size, ","); - absl::StrAppend(&debug_str, allocation.chunk().offset, ","); - absl::StrAppend(&debug_str, allocation.start_time(), ","); - absl::StrAppend(&debug_str, allocation.end_time(), "\n"); - } -} - -void MsaAlgorithm::DumpDebugStringsIfEnabled() const { - if (!options_.dump_fn) { - return; - } - options_.dump_fn("bufferinfo", buffer_info_str_); - options_.dump_fn("allocinfo", allocation_info_str_); - options_.dump_fn("scheduleinfo", instruction_schedule_str_); -} - -std::string MsaAlgorithm::OptionalRequiredMemoryAssignmentToString( - const std::optional& assignment) { - if (!assignment.has_value()) { - return "null"; - } - - if (assignment->offset) { - return absl::StrCat(MemorySpaceToString(assignment->memory_space), - " at offset ", assignment->offset->offset); - } - - return MemorySpaceToString(assignment->memory_space); -} - -absl::Status MsaAlgorithm::OptimizeMemoryBoundLoop(int loop_start_idx, - int loop_end_idx, - int loop_size) { - // The MemoryBoundLoopOptimizer works with a minimum of three unrolled loop - // iterations: previous, current, and next. So, we pick the second iteration - // out of the loop as the current iteration. - const int iteration_start_idx = loop_start_idx + loop_size; - const int iteration_end_idx = iteration_start_idx + loop_size; - - ABSL_ASSIGN_OR_RETURN(std::unique_ptr optimizer, - MemoryBoundLoopOptimizer::Create( - iteration_start_idx, iteration_end_idx, hlo_live_range_, - alias_analysis_, options_)); - optimizer->Optimize(); - - // Check if this unrolled loop is in a while loop. - const auto& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - std::vector callers = call_graph_->GetComputationCallers( - instruction_sequence[loop_start_idx]->parent()); - const bool is_in_while_loop = - callers.size() == 1 && callers.front()->opcode() == HloOpcode::kWhile; - - // Update the loop_optimized_allocations_map_ with the output of the - // optimizer. - for (MemoryBoundLoopOptimizer::LoopValue& value : optimizer->loop_values()) { - if (value.allocations.empty() || !value.IsAllocationTypeSupported()) { - continue; - } - loop_optimized_allocations_.push_back(std::move(value.allocations)); - const AllocationSequence& sequence = loop_optimized_allocations_.back(); - CHECK(!sequence.empty()); - VLOG(3) << " alloc: " << sequence.back()->ToString(); - for (const auto& allocation : sequence) { - // Check if the loop is in a while loop and the position needs to be - // allocated in the default memory. - const bool require_pos_in_default_space = - is_in_while_loop && - (allocation->memory_space() == MemorySpace::kDefault || - allocation->is_copy_allocation()); - for (const HloUse& use : allocation->uses()) { - const int64_t use_idx = - hlo_live_range_.instruction_schedule().at(use.instruction) - - iteration_start_idx; - CHECK_GE(use_idx, 0); - CHECK_LT(use_idx, loop_size); - for (int64_t i = loop_start_idx + use_idx; i <= loop_end_idx; - i += loop_size) { - HloInstruction* repeated_inst = instruction_sequence[i]; - CHECK_EQ(use.instruction->opcode(), repeated_inst->opcode()); - CHECK_EQ(use.instruction->operand_count(), - repeated_inst->operand_count()); - CHECK_LT(use.operand_number, repeated_inst->operand_count()); - HloUse repeated_use{repeated_inst, use.operand_number, - use.operand_index}; - loop_optimized_allocations_map_[repeated_use] = {use_idx, loop_size, - allocation.get()}; - VLOG(3) << " Setting optimized allocations map. Use: " - << repeated_use.ToString() << " idx: " << use_idx - << " allocation: " << allocation->ToString(); - if (require_pos_in_default_space) { - const HloValue& value = - alias_analysis_.dataflow_analysis().GetUniqueValueAt( - repeated_inst->operand(use.operand_number), - use.operand_index); - // If any of the positions is a parameter in a while loop, we add a - // required assignment in the default memory space. - for (const HloPosition& value_position : value.positions()) { - if (value_position.instruction->parent() == - repeated_inst->parent() && - value_position.instruction->opcode() == - HloOpcode::kParameter) { - // TODO(subhankarshah): Validate if adding to pending required - // assignments is intentional/correct, since we clear after - // every successful allocation. - AddRequiredAssignment( - value_position.instruction, value_position.index, - static_cast(MemorySpace::kDefault), - RequiredMemoryAssignment::Source:: - kLoopOptimizedParameterInWhileLoop, - nullptr, true); - break; - } - } - } - } - } - } - } - return absl::OkStatus(); -} - -namespace { -// A helper function to get the distance between a use and its producer (or -1 -// if producer is a gte, parameter or tuple). -std::function GetOperandDistanceFunction( - const HloLiveRange& hlo_live_range, const HloInstruction* use_inst) { - const int use_idx = hlo_live_range.instruction_schedule().at(use_inst); - return [&, use_idx](const HloInstruction* operand) -> int { - // We just use -1 for parameter, tuple, gte and constant instructions. We - // could make this "see through" the gtes if we get too many false - // positives. - if (operand->opcode() == HloOpcode::kParameter || - operand->opcode() == HloOpcode::kTuple || - operand->opcode() == HloOpcode::kGetTupleElement || - operand->opcode() == HloOpcode::kConstant) { - return -1; - } - return use_idx - hlo_live_range.instruction_schedule().at(operand); - }; -} - -// A helper function to check if the operand distances of two instructions -// are compatible. This assumes `a` is scheduled loop size candidate -// instructions before `b`. The operand distances are compatible if either -// distance is -1, or if they are the same, or if they are separated by loop -// size candidate. -bool AreOperandCandidatesCompatible(int loop_size_candidate, - absl::Span a_distances, - absl::Span b_distances) { - if (a_distances.size() != b_distances.size()) { - return false; - } - for (int i = 0; i < a_distances.size(); ++i) { - const int a_value = a_distances.at(i); - const int b_value = b_distances.at(i); - if (a_value != -1 && b_value != -1 && - a_value + loop_size_candidate != b_value && a_value != b_value) { - return false; - } - } - return true; -} -} // namespace - -void MsaAlgorithm::IdentifyAndOptimizeMemoryBoundLoops() { - absl::flat_hash_map fingerprint_schedule_map; - const auto& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - // The minimum and maximum loop sizes that we consider. - const int kMinLoopSize = 4; - const int kMaxLoopSize = 400; - int optimized_loop_idx = 0; - while (optimized_loop_idx < instruction_sequence.size()) { - // Iterate over the flattened instruction sequence. We first try to find a - // loop candidate where the fingerprint between two instructions matches by - // the loop size candidate. - int loop_size_candidate = -1; - int loop_start_idx = -1; - int loop_end_idx = -1; - for (; optimized_loop_idx < instruction_sequence.size(); - ++optimized_loop_idx) { - const HloInstruction* inst = instruction_sequence[optimized_loop_idx]; - auto fingerprint_it = fingerprint_map_.find(inst); - if (inst->opcode() != HloOpcode::kParameter && - inst->opcode() != HloOpcode::kTuple && - inst->opcode() != HloOpcode::kGetTupleElement && - fingerprint_it != fingerprint_map_.end()) { - // Find and the latest instruction with the same fingerprint as this. - auto fingerprint_schedule_it = - fingerprint_schedule_map.find(fingerprint_it->second); - if (fingerprint_schedule_it != fingerprint_schedule_map.end()) { - int distance = optimized_loop_idx - fingerprint_schedule_it->second; - if (distance >= kMinLoopSize && distance <= kMaxLoopSize) { - // We found two instructions with the same fingerprint. The distance - // between the two is the loop size candidate. - loop_size_candidate = distance; - // Update the fingerprint map with the current loop index so that if - // the loop size candidate doesn't find a valid loop, we can resume - // searching from this instruction. - fingerprint_schedule_map[fingerprint_it->second] = - optimized_loop_idx; - break; - } - } - fingerprint_schedule_map[fingerprint_it->second] = optimized_loop_idx; - } - - VLOG(3) << " " << optimized_loop_idx << ": " - << instruction_sequence[optimized_loop_idx]->parent()->name() - << " " << instruction_sequence[optimized_loop_idx]->name() - << " fingerprint: " - << (fingerprint_it == fingerprint_map_.end() - ? "none" - : std::to_string(fingerprint_it->second)); - } - VLOG(3) << "Loop size candidate: " << loop_size_candidate; - if (loop_size_candidate == -1) { - break; - } - - std::vector> operand_distances; - - // Scan the instructions with the candidate loop size. We try to calculate - // the size of the loop by finding the instructions that are loop size - // candidate apart, have the same fingerprint and compatible operand - // distances. We start scanning the candidate loop a few instructions - // earlier than the fingerprint identified in case the loop starts a bit - // earlier than the fingerprint logic. - const int kLoopScanHeadStart = 10; - for (int i = std::max( - 0, optimized_loop_idx - loop_size_candidate - kLoopScanHeadStart); - i < instruction_sequence.size(); ++i) { - const HloInstruction* inst = instruction_sequence[i]; - auto fingerprint_it = fingerprint_map_.find(inst); - auto ignore_op = [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kParameter || - instruction->opcode() == HloOpcode::kTuple || - instruction->opcode() == HloOpcode::kGetTupleElement; - }; - // We trigger this if statement until we find the start of the loop. - if (loop_start_idx == -1) { - if (i > optimized_loop_idx - loop_size_candidate) { - break; - } - if (ignore_op(inst) || fingerprint_it == fingerprint_map_.end()) { - continue; - } - if (i + loop_size_candidate >= instruction_sequence.size()) { - break; - } - const HloInstruction* candidate_inst = - instruction_sequence[i + loop_size_candidate]; - auto candidate_fingerprint_it = fingerprint_map_.find(candidate_inst); - if (ignore_op(candidate_inst) || - candidate_fingerprint_it == fingerprint_map_.end() || - fingerprint_it->second != candidate_fingerprint_it->second) { - // Fingerprint mismatch. - continue; - } - std::vector inst_operand_distances; - absl::c_transform(inst->operands(), - std::back_inserter(inst_operand_distances), - GetOperandDistanceFunction(hlo_live_range_, inst)); - std::vector candidate_inst_operand_distances; - absl::c_transform( - candidate_inst->operands(), - std::back_inserter(candidate_inst_operand_distances), - GetOperandDistanceFunction(hlo_live_range_, candidate_inst)); - VLOG(3) << "i : " << i << " " - << absl::StrJoin(inst_operand_distances, ", ") << " | " - << absl::StrJoin(candidate_inst_operand_distances, ", "); - if (!AreOperandCandidatesCompatible(loop_size_candidate, - inst_operand_distances, - candidate_inst_operand_distances)) { - // Operand distance mismatch. - continue; - } - // Found the start of the loop. - loop_start_idx = i; - } - if (inst->parent() != instruction_sequence[loop_start_idx]->parent()) { - VLOG(3) << "Mismatch (computation) at " << i << ": " - << inst->parent()->name() << " vs " - << instruction_sequence[loop_start_idx]->parent()->name(); - break; - } - operand_distances.push_back({}); - if (fingerprint_it == fingerprint_map_.end()) { - continue; - } - absl::c_transform(inst->operands(), - std::back_inserter(operand_distances.back()), - GetOperandDistanceFunction(hlo_live_range_, inst)); - if (i >= loop_start_idx + loop_size_candidate) { - // Verify that this still obeys the fingerprint and operand distance - // invariants. - const HloInstruction* prev_inst = - instruction_sequence[i - loop_size_candidate]; - auto prev_fingerprint_it = fingerprint_map_.find(prev_inst); - if (prev_fingerprint_it == fingerprint_map_.end()) { - break; - } - if (ignore_op(inst) || ignore_op(prev_inst)) { - if (inst->opcode() != prev_inst->opcode()) { - VLOG(3) << "Mismatch (opcode) at " << i << ", " - << (i - loop_size_candidate) << ": " << inst->opcode() - << " vs " << prev_inst->opcode(); - break; - } - if (inst->operand_count() != prev_inst->operand_count()) { - VLOG(3) << "Mismatch (# operands) at " << i << ", " - << (i - loop_size_candidate) << ": " - << inst->operand_count() << " vs " - << prev_inst->operand_count(); - break; - } - } - if (fingerprint_it->second != prev_fingerprint_it->second) { - VLOG(3) << "Mismatch (fp) at " << i << ", " - << (i - loop_size_candidate) << ": " << fingerprint_it->second - << " vs " << prev_fingerprint_it->second; - break; - } - if (!AreOperandCandidatesCompatible( - loop_size_candidate, - *(operand_distances.rbegin() + loop_size_candidate), - operand_distances.back())) { - VLOG(3) << "Mismatch (op) at " << i << ", " - << (i - loop_size_candidate) << ": " - << absl::StrJoin(operand_distances.back(), ", ") << " vs " - << absl::StrJoin( - *(operand_distances.rbegin() + loop_size_candidate), - ", "); - break; - } - } - loop_end_idx = i; - } - float num_iterations = 0; - if (loop_start_idx != -1) { - num_iterations = static_cast(loop_end_idx + 1 - loop_start_idx) / - loop_size_candidate; - } - VLOG(3) << "Loop start: " << loop_start_idx << " loop end: " << loop_end_idx - << " num iterations: " << num_iterations; - - optimized_loop_idx = std::max(optimized_loop_idx, loop_end_idx) + 1; - - if (num_iterations >= - options_.memory_bound_loop_optimizer_options.min_num_iterations()) { - VLOG(2) << "Found valid loop. Loop start: " << loop_start_idx - << " loop end: " << loop_end_idx - << " num iterations: " << num_iterations; - - CHECK_OK(OptimizeMemoryBoundLoop(loop_start_idx, loop_end_idx, - loop_size_candidate)); - } - } -} - -bool MsaAlgorithm::IsAsyncConversionCandidate( - const HloInstruction* instruction) const { - bool meets_special_preconditions = - IsAsyncConversionCopyCandidate(instruction) || - IsAsyncConversionSliceCandidate(instruction) == - AsyncConversionResult::kSuccess || - IsAsyncCustomFusionConversionCandidate(instruction) == - AsyncConversionResult::kSuccess; - if (!meets_special_preconditions) { - return false; - } - if (IsBlockPrefetchingEnabled()) { - return true; - } - - for (auto& operand : instruction->operands()) { - // TODO(b/374835319): relax the operand constraint to be able to cover - // nested sync data movement cases. - if (IsAsyncConversionCandidate(operand)) { - VLOG(4) << "The instruction is not considered to be replaced, because it " - "potentially has a replaceable operand."; - return false; - } - const HloValue& operand_value = alias_analysis_.dataflow_analysis() - .GetValueSet(operand) - .GetUniqueValue(); - if (!buffer_intervals_.at(&operand_value).need_allocation) { - VLOG(4) - << "The instruction is not considered to be replaced, because its " - "operand value doesn't need an allocation."; - return false; - } - } - - const HloValue& value = alias_analysis_.dataflow_analysis() - .GetValueSet(instruction) - .GetUniqueValue(); - if (!buffer_intervals_.at(&value).need_allocation) { - VLOG(4) << "The instruction is not considered to be replaced, because its " - "output doesn't need an allocation and it might be too late to " - "replace this instruction."; - return false; - } - if (value.IsRootOf(instruction->parent())) { - VLOG(4) << "The instruction is not considered to be replaced, because its " - "output value is in the root of the computation."; - return false; - } - if (IsValueFinalized(&value)) { - VLOG(4) << "The instruction is not considered to be replaced, because its " - "output value is in the finalized values."; - return false; - } - if (buffer_intervals_.at(&value).size > available_heap_size()) { - VLOG(4) << "The instruction is not considered to be replaced, because its " - "output value is too large to fit in the heap."; - return false; - } - // This check is here only because we skip processing the values that are not - // allowed in alternate memory. - if (!MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory( - buffer_intervals_.at(&value), options_.alternate_memory_space)) { - VLOG(4) << "The instruction is not considered to be replaced, because its " - "output value is not allowed in alternate memory."; - return false; - } - - for (const HloInstruction* user : instruction->users()) { - if (HloDataflowAnalysis::IsAsynchronousOperationStart(user->opcode())) { - VLOG(4) << "The instruction is not considered to be replaced, because " - "its used by an async start operation that might require " - "contiguous allocation."; - return false; - } - } - - return true; -} - -bool MsaAlgorithm::IsAsyncConversionCopyCandidate( - const HloInstruction* instruction) const { - if (!options_.enable_sync_copy_replacement) { - return false; - } - if (instruction->opcode() != HloOpcode::kCopy) { - return false; - } - if (failed_async_conversions_.contains(instruction)) { - return false; - } - if (instruction->IsRoot()) { - // Root copy is not replaceable with current implementation, because the - // instruction has no uses - return false; - } - if (instruction->operand(0)->shape() != instruction->shape()) { - VLOG(5) << "Sync copy " << instruction->ToShortString() - << " is not replaceable, because the operand and output shapes do " - "not match. This could be because the copy is changing the " - "layout. operand(0) shape: " - << instruction->operand(0)->shape().ToString(/*print_layout=*/true) - << ", output shape: " - << instruction->shape().ToString(/*print_layout=*/true); - return false; - } - if (instruction->shape().layout().memory_space() != - static_cast(MemorySpace::kDefault) || - instruction->operand(0)->shape().layout().memory_space() != - static_cast(MemorySpace::kDefault)) { - VLOG(5) << "Sync copy " << instruction->ToShortString() - << " is not replaceable, because the operand or output have an " - "initial assignment."; - return false; - } - return true; -} - -MsaAlgorithm::AsyncConversionResult -MsaAlgorithm::IsAsyncCustomFusionConversionCandidate( - const HloInstruction* instruction) const { - if (!options_.enable_sync_custom_fusion_replacement) { - return AsyncConversionResult::kFeatureNotEnabled; - } - auto it = failed_async_conversions_.find(instruction); - if (it != failed_async_conversions_.end()) { - return it->second; - } - if (!instruction->IsCustomFusion() || - !options_.custom_fusion_block_prefetch_operand_index_fn(instruction) - .has_value()) { - return AsyncConversionResult::kFailedPrecondition; - } - return AsyncConversionResult::kSuccess; -} - -MsaAlgorithm::AsyncConversionResult -MsaAlgorithm::IsAsyncConversionSliceCandidate( - const HloInstruction* instruction) const { - if (!options_.enable_sync_slice_replacement) { - return AsyncConversionResult::kFeatureNotEnabled; - } - if (failed_async_conversions_.contains(instruction)) { - return failed_async_conversions_.at(instruction); - } - if (!IsSliceLikeInstruction(instruction)) { - return AsyncConversionResult::kFailedPrecondition; - } - - if (instruction->IsRoot()) { - // Root slice is not replaceable with current implementation, because the - // instruction has no uses - return AsyncConversionResult::kFailedPrecondition; - } - - if (!options_.is_async_slice_implemented_fn(instruction)) { - VLOG(4) << "The sync slice is not considered to be replaced, because the " - "async version is not implemented for " - << instruction->ToShortString(); - return AsyncConversionResult::kFailedPrecondition; - } - - if (instruction->shape().layout().memory_space() != - static_cast(MemorySpace::kDefault) || - instruction->operand(0)->shape().layout().memory_space() != - static_cast(MemorySpace::kDefault)) { - VLOG(4) << "Sync slice " << instruction->ToShortString() - << " is not replaceable, because the operand or output have an " - "initial assignment."; - return AsyncConversionResult::kFailedPrecondition; - } - if (instruction->shape().element_type() != - instruction->operand(0)->shape().element_type()) { - VLOG(4) << "Sync slice " << instruction->ToShortString() - << " is not replaceable because the operand and output have " - "different element types."; - return AsyncConversionResult::kFailedPrecondition; - } - return AsyncConversionResult::kSuccess; -} - -std::vector MsaAlgorithm::GenerateJointProcessedValues( - const HloValue* entrance_value) { - std::vector worklist = {entrance_value}; - if (!IsBlockPrefetchingEnabled() && - (options_.enable_sync_copy_replacement || - options_.enable_sync_slice_replacement)) { - // Adds the HloValue that is related to a given instruction to the worklist - auto add_to_worklist = [&](const HloInstruction* inst) { - const HloValue& next_value = alias_analysis_.dataflow_analysis() - .GetValueSet(inst) - .GetUniqueValue(); - if (absl::c_find(worklist, &next_value) == worklist.end()) { - worklist.push_back(&next_value); - } - }; - - for (size_t idx = 0; idx < worklist.size(); ++idx) { - const HloValue* value = worklist.at(idx); - // Values that are related to the current value through a sync copy use - // are added to the worklist. - for (const auto& use : value->GetUses()) { - if (IsAsyncConversionCandidate(use.instruction)) { - add_to_worklist(use.instruction); - } - } - // Expand the worklist to include values that connect to the current - // value as sync copy operands, if any. - HloInstruction* defining_instruction = value->instruction(); - if (IsAsyncConversionCandidate(defining_instruction)) { - // The first operand of slice like instruction (slice or dynamic-slice) - // is the defining instruction. - add_to_worklist(defining_instruction->operand(0)); - } - } - // We're sensitive to the order of the worklist. - absl::c_stable_sort(worklist, [&](const HloValue* a, const HloValue* b) { - return std::make_pair(buffer_intervals_.at(a).start, - buffer_intervals_.at(a).end) < - std::make_pair(buffer_intervals_.at(b).start, - buffer_intervals_.at(b).end); - }); - UpdateSyncDataMovementCandidatesForJointProcessedValues(worklist); - } - return worklist; -} - -void MsaAlgorithm::UpdateSyncDataMovementCandidatesForJointProcessedValues( - const std::vector& joint_processed_values) { - absl::flat_hash_set replaceable_sync_instructions; - for (const HloValue* value : joint_processed_values) { - for (const auto& use : value->GetUses()) { - if (IsAsyncConversionCandidate(use.instruction)) { - replaceable_sync_instructions.insert(use.instruction); - } - } - HloInstruction* inst = value->instruction(); - if (IsAsyncConversionCandidate(inst)) { - replaceable_sync_instructions.insert(inst); - } - } - - sorted_async_conversion_candidates_.clear(); - sorted_async_conversion_candidates_.insert( - sorted_async_conversion_candidates_.end(), - replaceable_sync_instructions.begin(), - replaceable_sync_instructions.end()); - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - absl::c_stable_sort(sorted_async_conversion_candidates_, - [&instruction_schedule](const HloInstruction* a, - const HloInstruction* b) { - return instruction_schedule.at(a) < - instruction_schedule.at(b); - }); - VLOG(3) << "Sorted pending replaceable copies: "; - if (sorted_async_conversion_candidates_.empty()) { - VLOG(3) << " --Empty--"; - } - for (size_t idx = 0; idx < sorted_async_conversion_candidates_.size(); - ++idx) { - VLOG(3) << " " << idx + 1 << "/" - << sorted_async_conversion_candidates_.size() << ") " - << sorted_async_conversion_candidates_.at(idx)->ToShortString(); - } -} - -namespace { - -void NicePrintAllocationValues( - const std::vector& allocation_values, int log_level) { - if (VLOG_IS_ON(log_level)) { - VLOG(log_level) << "Joint-processed allocation values: "; - if (allocation_values.empty()) { - VLOG(log_level) << " --Empty--"; - } - for (int i = 0; i < allocation_values.size(); i++) { - VLOG(log_level) << " " << i + 1 << "/" << allocation_values.size() - << ") " << allocation_values.at(i).ToShortString(); - } - } -} - -} // namespace - -void MsaAlgorithm::ColorColocatedIntervalsToAlternate( - const std::vector& colocated_intervals) { - for (const MsaBufferInterval* colocated_interval : colocated_intervals) { - const HloValue* value = colocated_interval->buffer; - // Color all of the aliased reserved buffers here because reserved - // alternate memory allocations will not have an entry in preset - // allocations that is normally used for coloring. - for (auto& position : value->positions()) { - VLOG(4) << "Coloring " << position.ToString(); - Shape* shape = ShapeUtil::GetMutableSubshape( - position.instruction->mutable_shape(), position.index); - CHECK(shape->IsArray()) - << "Coloring a shape that is not an array: " << position.ToString(); - shape->mutable_layout()->set_memory_space( - options_.alternate_memory_space); - } - } -} - -void MsaAlgorithm::CreateAllocationValuesForJointProcessedValues( - JointAllocationProposal& proposal) { - for (const HloValue* interval_hlo_value : proposal.values) { - const MsaBufferInterval& interval = - buffer_intervals_.at(interval_hlo_value); - - if (IsValueFinalized(interval_hlo_value)) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because it is already processed."; - continue; - } - - if (!interval.need_allocation) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because it doesn't need an allocation."; - continue; - } - - if (!MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory( - interval, options_.alternate_memory_space)) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because it is not allowed in the alternate memory."; - continue; - } - - HloInstruction* inst = interval.buffer->instruction(); - HloModule* module = inst->GetModule(); - - // Don't intra-program prefetch a cross program prefetch - auto cross_program_prefetches = module->CrossProgramPrefetches(); - if (inst->opcode() == HloOpcode::kParameter && - absl::c_find_if(cross_program_prefetches, [&](auto& info) { - return info.parameter == inst->parameter_number() && - info.index == interval.buffer->index(); - }) != module->CrossProgramPrefetches().end()) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because it is cross-program prefetched."; - continue; - } - - if (!options_.IsOpSpanExposureEnabled() && - interval.size > available_heap_size()) { - const HloInstruction* defining_instruction = - interval.buffer->instruction(); - auto may_be_replaced_by_slice_fn = [this](const HloInstruction* user) { - return IsInstructionPendingReplacements(user) && - IsSliceLikeInstruction(user); - }; - bool may_be_replaced_by_slice = std::any_of( - defining_instruction->users().begin(), - defining_instruction->users().end(), may_be_replaced_by_slice_fn); - - if (!may_be_replaced_by_slice) { - if (!MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, - interval.buffer->defining_position(), interval.size)) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because the buffer is larger than the heap size."; - continue; - } - } - } - - std::vector colocated_intervals = - GetSortedColocatedIntervals(interval); - if (AreIntervalsReservedInAlternateMemory(colocated_intervals)) { - VLOG(3) << "Interval " << interval.buffer->ToShortString() - << " is reserved in the alternate memory."; - ColorColocatedIntervalsToAlternate(colocated_intervals); - continue; - } - - if (!ConsumeFuel("memory_space_assignment", [&] { - return absl::StrCat("Ran out of fuel at buffer: ", - colocated_intervals[0]->buffer->ToShortString()); - })) { - continue; - } - - if (options_.dump_fn != nullptr || VLOG_IS_ON(3)) { - // Only fill buffer_info_str_ if needed. - AppendBufferInfoDebugString(interval, &buffer_info_str_); - } - - CreateAllocationValuesFromColocatedIntervals(colocated_intervals, - proposal.allocation_values); - proposal.colocated_intervals.push_back(colocated_intervals); - } - // Order allocation values so that when read and write sequences are - // different, we're sure the allocation value corresponding to the read - // sequence is processed before the written sequence's allocation value. - // We move the allocation values that have their defining instruction in the - // sync conversion list to the end to be processed last. - std::stable_partition(proposal.allocation_values.begin(), - proposal.allocation_values.end(), - [this](AllocationValue& allocation_value) { - return !IsInstructionPendingReplacements( - allocation_value.position().instruction); - }); - - NicePrintAllocationValues(proposal.allocation_values, /*log_level=*/3); -} - -MsaAlgorithm::JointAllocationProposal MsaAlgorithm::GetJointProposal( - MsaBufferInterval& interval) { - JointAllocationProposal proposal; - proposal.values = GenerateJointProcessedValues(interval.buffer); - if (VLOG_IS_ON(3)) { - VLOG(3) << "Joint-processed values for " << interval.buffer->ToShortString() - << ": "; - for (size_t idx = 0; idx < proposal.values.size(); ++idx) { - const HloValue* hlo_value = proposal.values.at(idx); - VLOG(3) << " " << idx + 1 << "/" << proposal.values.size() << ") " - << hlo_value->ToShortString(); - } - } - - CreateAllocationValuesForJointProcessedValues(proposal); - return proposal; -} - -bool MsaAlgorithm::RepackAllocationsIncludeConvertedSyncMemOp() { - for (RepackAllocationBlock& allocation_block : repack_allocation_blocks_) { - if (allocation_block.allocation->is_copy_allocation()) { - if (dynamic_cast(allocation_block.allocation) - ->sync_mem_op()) { - return true; - } - } - if (allocation_block.allocation->is_sliced_copy_allocation()) { - if (dynamic_cast(allocation_block.allocation) - ->sync_mem_op()) { - return true; - } - } - } - return false; -} - -namespace { - -// Fixes the AllocationSequence after post-allocation transformation: -// 1. Remove the allocations with to_be_removed instructions as the defining -// positions. -// 2. Update the vector of uses for all allocations according to the -// update_use_map. -// Note that to_be_removed instructions will later be removed from the module -// during SimplifyGraph() call in memory_space_assignment.cc -void FixAllocationSequenceAfterPostAllocationTransformation( - AllocationSequence* allocations, - const PostAllocationTransformationUpdate& transformation_info) { - VLOG(3) << "Fixing AllocationSequence after post-allocation transformation"; - - // (1) - allocations->erase( - std::remove_if( - allocations->begin(), allocations->end(), - [transformation_info](const std::unique_ptr& allocation) { - return absl::c_find(transformation_info.to_be_removed, - allocation->defining_position().instruction) != - transformation_info.to_be_removed.end(); - }), - allocations->end()); - - // (2) - for (auto& allocation : *allocations) { - std::vector uses_to_update; - for (const HloUse& use : allocation->uses()) { - for (const auto& [old_use, new_use] : - transformation_info.update_use_map) { - if (use == old_use) { - uses_to_update.push_back(old_use); - break; // found the use, no need to keep searching update_use_map - } - } - } - - // Perform update uses - if (!uses_to_update.empty()) { - for (const HloUse& old_use : uses_to_update) { - const HloUse& new_use = transformation_info.update_use_map.at(old_use); - allocation->RemoveUse(old_use); - allocation->AddUse(new_use); - } - } - } -} - -// Verifies that the operands_in_alternate_memory_map is consistent with the -// allocations in the AllocationSequence. -bool VerifyOperandsInAlternateMemoryMap( - const AllocationSequence* allocations, - const absl::flat_hash_map>>& - operands_in_alternate_memory_map) { - absl::flat_hash_map>> - reference_map; - for (const std::unique_ptr& - allocation : *allocations) { - if (allocation->is_in_alternate_mem()) { - for (const HloUse& use : allocation->uses()) { - reference_map[use.instruction].insert( - std::make_pair(use.operand_number, use.operand_index)); - } - } - } - return reference_map == operands_in_alternate_memory_map; -} - -// GetAsyncCopyElapsed with a default value. -float CopyResourceForShape(const Options& options, const Shape& shape) { - return options.cost_analysis - ? options.cost_analysis->GetAsyncCopyElapsed( - options.cost_analysis->GetShapeSizeBytes(shape)) - : 0.1; -} - -absl::StatusOr GetMemorySpaceEnum(const int64_t memory_space, - const Options& options) { - if (memory_space == options.alternate_memory_space) { - return MemorySpace::kAlternate; - } - if (memory_space == options.default_memory_space) { - return MemorySpace::kDefault; - } - return absl::InvalidArgumentError( - absl::StrCat("Invalid memory space: ", memory_space)); -} - -// Sorts and merges time intervals. If `allow_overlapping_intervals` is true, -// it logs a warning on overlap instead of returning an error. -absl::StatusOr> SortAndMergeTimeIntervals( - std::vector all_ranges, bool allow_overlapping_intervals) { - if (all_ranges.empty()) { - return std::vector(); - } - absl::c_sort(all_ranges); - std::vector merged; - for (const auto& r : all_ranges) { - if (merged.empty()) { - merged.push_back(r); - } else if (merged.back().inclusive_end_time < r.inclusive_start_time) { - merged.push_back(r); - } else if (allow_overlapping_intervals) { - if (merged.back().inclusive_end_time > r.inclusive_start_time) { - // This case should never happen for a valid buffer coloring. - // Positions may have overlapping live ranges at only at aliased - // boundaries. Which would result in equal end/start times. - VLOG(1) << "WARNING: Overlapping live ranges detected [" - << merged.back().inclusive_start_time << ", " - << merged.back().inclusive_end_time << "] and [" - << r.inclusive_start_time << ", " << r.inclusive_end_time - << "], possible buffer corruption bug.\n"; - } - merged.back().inclusive_end_time = - std::max(merged.back().inclusive_end_time, r.inclusive_end_time); - } else { - return absl::InvalidArgumentError(absl::StrCat( - "Overlapping live ranges detected [", - merged.back().inclusive_start_time, ", ", - merged.back().inclusive_end_time, "] and [", r.inclusive_start_time, - ", ", r.inclusive_end_time, "]")); - } - } - return merged; -} - -// Appends a time interval to `intervals` for a coloring time. If the coloring -// time falls within a contiguous live range, the entire contiguous live range -// is added to `intervals`, else a time interval of one time unit, i.e. -// [time_of_coloring, time_of_coloring] is added. -void AddIntervalForColoringTime( - int64_t time_of_coloring, std::vector& intervals, - std::vector& contiguous_live_ranges) { - TimeInterval new_interval = {time_of_coloring, time_of_coloring}; - // Find the first interval that ends at or after the coloring time. - auto it = std::lower_bound(contiguous_live_ranges.begin(), - contiguous_live_ranges.end(), time_of_coloring, - [](const TimeInterval& interval, int64_t t) { - return interval.inclusive_end_time < t; - }); - if (it != contiguous_live_ranges.end() && - it->inclusive_start_time <= time_of_coloring) { - new_interval = *it; - } - intervals.push_back(new_interval); -} - -} // namespace - -absl::StatusOr> -MsaAlgorithm::GetContiguousLiveRangesForBuffer(const HloBuffer* buffer) const { - absl::flat_hash_map> - contiguous_positions_to_uses; - for (const HloValue* value : buffer->values()) { - if (!value->shape().IsArray()) { - continue; - } - for (const HloUse& use : value->GetUses()) { - HloInstruction* operand = - use.instruction->mutable_operand(use.operand_number); - ShapeIndex operand_index = use.operand_index; - HloPosition source_position = - GetNonTrivialSourcePosition(HloPosition{operand, operand_index}); - bool is_pipelined_async = - (has_async_pipelined_while_loops_ && - IsBufferAliasedToAsyncPipelinedWhileLoop(value) && - (IsAsyncOperationStateDefinition(source_position.instruction) || - IsAsyncPipelinedWhilePosition(source_position) || - IsAsyncOperationStateUse(use))); - if (options_.position_requires_contiguous_allocation_fn( - source_position) || - is_pipelined_async) { - VLOG(3) << "Adding use " << use.ToString() << " to contiguous position " - << source_position.ToString(); - contiguous_positions_to_uses[source_position].push_back(use); - } - } - for (const HloPosition& position : value->positions()) { - bool is_pipelined_async_pos = - (has_async_pipelined_while_loops_ && - IsBufferAliasedToAsyncPipelinedWhileLoop(value) && - (IsAsyncOperationStateDefinition(position.instruction) || - IsAsyncPipelinedWhilePosition(position))); - if (options_.position_requires_contiguous_allocation_fn(position) || - is_pipelined_async_pos) { - if (!(contiguous_positions_to_uses.contains(position))) { - LOG(WARNING) << "Position " << position.ToString() - << " is required to be contiguous but has no uses, " - "this should not happen."; - contiguous_positions_to_uses[position] = {}; - } - } - } - } - std::vector contiguous_live_ranges; - - for (const auto& [position, uses] : contiguous_positions_to_uses) { - int64_t start_time = - hlo_live_range_.instruction_schedule().at(position.instruction); - int64_t end_time = start_time; - if (uses.empty() && - position.instruction->opcode() == HloOpcode::kParameter && - position.instruction->parent() != nullptr && - position.instruction->parent()->root_instruction() != nullptr && - IsWhileBodyOrConditionComputation(position.instruction->parent())) { - HloInstruction* while_caller = GetWhileForBodyRoot( - position.instruction->parent()->root_instruction()); - if (while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller)) { - end_time = std::max( - end_time, hlo_live_range_.instruction_schedule().at( - position.instruction->parent()->root_instruction())); - } - } - for (const HloUse& use : uses) { - end_time = std::max(end_time, GetCorrectedUseTime(use)); - } - contiguous_live_ranges.push_back(TimeInterval{start_time, end_time}); - } - VLOG(3) << "Found " << contiguous_live_ranges.size() - << " contiguous live ranges for buffer " << buffer->ToString(); - return SortAndMergeTimeIntervals(contiguous_live_ranges, - /*allow_overlapping_intervals=*/true); -} - -absl::StatusOr> -MsaAlgorithm::GetAltMemoryColoredIntervalsForBuffer( - const HloBuffer* buffer, - const std::vector& buffer_colorings) { - std::vector default_mem_intervals; - std::vector alternate_mem_intervals; - ABSL_ASSIGN_OR_RETURN(std::vector contiguous_live_ranges, - GetContiguousLiveRangesForBuffer(buffer)); - - auto disallow_async_conversion_if_conversion_candidate = - [&](const HloInstruction* inst) { - if (IsAsyncConversionCandidate(inst)) { - failed_async_conversions_[inst] = - AsyncConversionResult::kAsyncConversionNotAllowedForColoredBuffer; - } - }; - - for (const auto& buffer_coloring : buffer_colorings) { - ABSL_ASSIGN_OR_RETURN( - const MemorySpace memory_space_enum, - GetMemorySpaceEnum(buffer_coloring.memory_space, options_)); - HloInstruction* coloring_site; - int64_t time_of_coloring; - if (std::holds_alternative( - buffer_coloring.buffer_position_or_use)) { - HloUse use = std::get(buffer_coloring.buffer_position_or_use); - coloring_site = use.instruction; - time_of_coloring = GetCorrectedUseTime(use); - HloInstruction* operand = - use.instruction->mutable_operand(use.operand_number); - ShapeIndex operand_index = use.operand_index; - HloInstruction* non_trivial_source_instruction = - GetNonTrivialSourcePosition(HloPosition{operand, operand_index}) - .instruction; - disallow_async_conversion_if_conversion_candidate( - non_trivial_source_instruction); - } else { - coloring_site = - std::get(buffer_coloring.buffer_position_or_use) - .instruction; - time_of_coloring = - hlo_live_range_.instruction_schedule().at(coloring_site); - } - - // Disallow async conversion if the coloring site is a candidate for async - // conversion. - disallow_async_conversion_if_conversion_candidate(coloring_site); - - if (coloring_site->opcode() == HloOpcode::kConditional || - coloring_site->opcode() == HloOpcode::kWhile) { - // TODO(subhankarshah): Add support for reserving and releasing chunks - // with disjoint live ranges for proper handling of while and conditional - // buffer coloring. - LOG(WARNING) << "Coloring while and conditional input/output is best " - "effort, might fail." - << coloring_site->ToString(); - } - std::vector& target_intervals = - (memory_space_enum == MemorySpace::kDefault) ? default_mem_intervals - : alternate_mem_intervals; - AddIntervalForColoringTime(time_of_coloring, target_intervals, - contiguous_live_ranges); - } - ABSL_ASSIGN_OR_RETURN( - std::vector merged_default_mem_intervals, - SortAndMergeTimeIntervals(default_mem_intervals, - /*allow_overlapping_intervals=*/true)); - ABSL_ASSIGN_OR_RETURN( - std::vector merged_alternate_mem_intervals, - SortAndMergeTimeIntervals(alternate_mem_intervals, - /*allow_overlapping_intervals=*/true)); - - VLOG(3) << "Buffer " << buffer->ToString() << " has " - << merged_default_mem_intervals.size() - << " default memory intervals and " - << merged_alternate_mem_intervals.size() - << " alternate memory intervals."; - - ABSL_RETURN_IF_ERROR(CheckForConflictingColoringRequirements( - buffer, merged_default_mem_intervals, merged_alternate_mem_intervals)); - - // Update the default memory coloring requirements for the buffer and return - // the intervals colored in alternate memory space. - default_memory_coloring_requirements_[buffer] = merged_default_mem_intervals; - return merged_alternate_mem_intervals; -} - -absl::Status MsaAlgorithm::CheckForConflictingColoringRequirements( - const HloBuffer* buffer, - const std::vector& default_mem_intervals, - const std::vector& alternate_mem_intervals) const { - std::vector overlap_check; - overlap_check.insert(overlap_check.end(), default_mem_intervals.begin(), - default_mem_intervals.end()); - overlap_check.insert(overlap_check.end(), alternate_mem_intervals.begin(), - alternate_mem_intervals.end()); - - if (!SortAndMergeTimeIntervals(overlap_check, - /*allow_overlapping_intervals=*/false) - .ok()) { - return absl::FailedPreconditionError( - "Conflicting memory space colorings exist for buffer: " + - buffer->ToString()); - } - return absl::OkStatus(); -} - -absl::linked_hash_map> -MsaAlgorithm::GetHloBufferToColoringsMap() const { - absl::linked_hash_map> - buffer_to_colorings; - for (const auto& buffer_coloring : options_.buffer_colorings) { - bool is_position = std::holds_alternative( - buffer_coloring.buffer_position_or_use); - if (is_position) { - HloPosition position = - std::get(buffer_coloring.buffer_position_or_use); - // Skip non-array shapes (e.g., tuples and tokens) since they do not - // consume physical alternate memory storage. - if (!position.shape().IsArray()) { - continue; - } - const HloBuffer& hlo_buffer = alias_analysis_.GetUniqueBufferAt( - position.instruction, position.index); - buffer_to_colorings[&hlo_buffer].push_back(buffer_coloring); - } else { - HloUse use = std::get(buffer_coloring.buffer_position_or_use); - HloInstruction* operand = - use.instruction->mutable_operand(use.operand_number); - ShapeIndex operand_index = use.operand_index; - const HloBuffer& hlo_buffer = - alias_analysis_.GetUniqueBufferAt(operand, operand_index); - // Skip non-array shapes (e.g., tuples and tokens) since they do not - // consume physical alternate memory storage. - if (!hlo_buffer.values().empty() && - !hlo_buffer.values().front()->shape().IsArray()) { - continue; - } - buffer_to_colorings[&hlo_buffer].push_back(buffer_coloring); - } - } - return buffer_to_colorings; -} - -absl::Status MsaAlgorithm::ProcessColoredBuffers() { - absl::linked_hash_map> - buffer_to_colorings = GetHloBufferToColoringsMap(); - - for (const auto& [buffer, buffer_colorings] : buffer_to_colorings) { - VLOG(3) << "Processing colored buffer " << buffer->ToString() << " with " - << buffer_colorings.size() << " colorings."; - // Get the size of the buffer. - CHECK(!buffer->values().empty()) - << "Colored buffer has no values: " << buffer->ToString(); - const HloValue* first_value = buffer->values().front(); - // Skip non-array shapes (e.g., tuples and tokens) since they do not - // consume physical alternate memory storage. - if (!first_value->shape().IsArray()) { - continue; - } - auto buffer_interval_it = buffer_intervals_.find(first_value); - if (buffer_interval_it == buffer_intervals_.end() && - first_value->defining_instruction()->opcode() == HloOpcode::kConstant) { - // Skip coloring for constant buffers. - continue; - } - CHECK(buffer_interval_it != buffer_intervals_.end()) - << "Colored buffer value not found in buffer intervals: " - << first_value->ToString() << ", buffer " << buffer->ToString() - << " skipping coloring."; - - int64_t buffer_size = buffer_interval_it->second.size; - - // Mark the instruction as ineligible for async conversion if it is a - // candidate for async conversion. - if (IsAsyncConversionCandidate( - first_value->defining_position().instruction)) { - failed_async_conversions_[first_value->defining_position().instruction] = - AsyncConversionResult::kAsyncConversionNotAllowedForColoredBuffer; - } - - // Find the intervals that are colored in alternate memory space and reserve - // the memory for them in the alternate memory space. - std::vector alternate_memory_space_intervals; - ABSL_ASSIGN_OR_RETURN( - alternate_memory_space_intervals, - GetAltMemoryColoredIntervalsForBuffer(buffer, buffer_colorings)); - for (const auto& alt_memory_space_interval : - alternate_memory_space_intervals) { - MsaBufferInterval interval = { - /*buffer=*/nullptr, - /*size=*/buffer_size, - /*start=*/alt_memory_space_interval.inclusive_start_time, - /*end=*/alt_memory_space_interval.inclusive_end_time, - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk chunk_candidate = FindChunkCandidate(interval); - VLOG(3) << "Found chunk candidate for buffer " << buffer->ToString() - << " interval " << alt_memory_space_interval.ToString() << " : " - << chunk_candidate.ToString(); - if (chunk_candidate.chunk_end() > options_.max_size_in_bytes) { - return FailedPrecondition( - "%s", absl::StrCat( - "Too many buffers are colored in the alternate " - "memory. Could not reserve alternate memory for buffer: ", - buffer->ToString(), " with size: ", buffer_size, - " for colored interval: ", interval.ToString())); - } - std::vector>& reserved_allocations = - reserved_allocations_for_alt_mem_colorings_[buffer]; - reserved_allocations.push_back(std::make_unique( - first_value->defining_position(), chunk_candidate, interval.start, - interval.end)); - AddToPendingChunks(interval, chunk_candidate); - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - interval.start, interval.end, chunk_candidate.size, - chunk_candidate.offset, reserved_allocations.back().get())); - repack_allocation_blocks_.back().next_colocated = - &(repack_allocation_blocks_.back()); - } - } - ClearPendingChunks(); - return absl::OkStatus(); -} - -int64_t MsaAlgorithm::MaxScopedMemorySize() { - const std::vector& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - int64_t max_reserved_scoped_memory = 0; - for (int64_t i = 0; i < instruction_sequence.size(); ++i) { - HloInstruction* instruction = instruction_sequence[i]; - int64_t reserved_scoped_memory = - std::min(options_.reserved_scoped_memory_fn( - instruction, /*operands_in_alternate_memory=*/{}, - /*outputs_in_alternate_memory=*/{}), - options_.max_size_in_bytes); - max_reserved_scoped_memory = - std::max(max_reserved_scoped_memory, reserved_scoped_memory); - } - VLOG(1) << "Max reserved scoped memory: " << max_reserved_scoped_memory; - return max_reserved_scoped_memory; -} - -std::optional MsaAlgorithm::EarliestBlockPrefetchStartTime( - int64_t previous_block_start_time, int64_t definition_time, - int64_t first_use_time, int64_t last_use_time, int64_t buffer_size, - int64_t block_prefetching_limit_bytes, - int64_t max_in_flight_prefetches_allowed, - std::vector& copy_done_schedule_before_times, - std::vector& block_prefetch_allocation_end_times) { - auto can_find_chunk_within_limit = - [&](int64_t start_time, int64_t end_time, int64_t buffer_size, - int64_t block_prefetching_limit_bytes) -> bool { - MsaBufferInterval interval = MsaBufferInterval{/*buffer=*/nullptr, - /*size=*/buffer_size, - /*start=*/start_time, - /*end=*/end_time, - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk chunk_candidate = FindChunkCandidate(interval); - return chunk_candidate.chunk_end() <= block_prefetching_limit_bytes; - }; - CHECK_LE(previous_block_start_time, first_use_time); - int64_t earliest_start_time_candidate = - std::max(definition_time, previous_block_start_time); - int64_t total_prefetches = copy_done_schedule_before_times.size(); - if (total_prefetches >= max_in_flight_prefetches_allowed) { - earliest_start_time_candidate = std::max( - earliest_start_time_candidate, - copy_done_schedule_before_times[total_prefetches - - max_in_flight_prefetches_allowed]); - } - if (can_find_chunk_within_limit(earliest_start_time_candidate, last_use_time, - buffer_size, block_prefetching_limit_bytes)) { - return earliest_start_time_candidate; - } - - // We couldn't find a chunk using earliest_start_time_candidate as our start - // So, we need to search for a valid start time in - // [earliest_start_time_candidate + 1, first_use_time], at the boundaries of - // the previous end times, which free up memory. - - // `it_begin` will point to the first element in the prefetch_end_times list - // that is strictly greater than (earliest_start_time_candidate - 1). Note, in - // the loop below, we add one to start time candidates. - auto it_begin = std::upper_bound(block_prefetch_allocation_end_times.begin(), - block_prefetch_allocation_end_times.end(), - earliest_start_time_candidate - 1); - - // `it_end` will point to the first element in the prefetch_end_times list - // that is strictly greater than (first_use_time - 1). - auto it_end = std::upper_bound(block_prefetch_allocation_end_times.begin(), - block_prefetch_allocation_end_times.end(), - first_use_time - 1); - for (auto it = it_begin; it != it_end; ++it) { - int64_t start_time = *it + 1; - if (can_find_chunk_within_limit(start_time, last_use_time, buffer_size, - block_prefetching_limit_bytes)) { - return start_time; - } - } - return std::nullopt; -} - -std::optional MsaAlgorithm::GetLatestSourceOperandScheduleTime( - HloInstruction* instruction, - const absl::flat_hash_map& - instruction_schedule) const { - int64_t latest_source_operand_time = std::numeric_limits::min(); - for (HloInstruction* operand : instruction->mutable_operands()) { - HloPosition non_trivial_source = - GetNonTrivialSourcePosition(HloPosition{operand, {}}); - // If the source instruction is kConstant, it can be freely moved. - // If there are other instructions that can be freely moved, - // we should consider them as well. This is an optimization and not - // required for correctness. - if (non_trivial_source.instruction->opcode() == HloOpcode::kConstant) { - continue; - } - // The previous passes should have done their best to hoist all - // operands of the async conversion candidates as early as possible. - latest_source_operand_time = std::max( - latest_source_operand_time, - FindOrDefault(instruction_schedule, non_trivial_source.instruction, - latest_source_operand_time)); - } - if (latest_source_operand_time == std::numeric_limits::min()) { - return std::nullopt; - } - return latest_source_operand_time; -} - -namespace { - -// Computes each value's [first_use_time, last_use_time] interval. When -// `view_color` is set, a use by a view colored instruction extends the last -// use time through the view's transitive readers (see -// ViewExtendedTransitiveUseTime): the view carries no storage of its own, so -// the value's buffer must stay reserved while any consumer still reads it -// through the view. -absl::flat_hash_map GetUseIntervals( - const std::vector& values, - const absl::flat_hash_map& - instruction_schedule, - std::optional view_color) { - absl::flat_hash_map value_to_use_intervals; - for (const HloValue* value : values) { - UseInterval& use_interval = value_to_use_intervals[value]; - use_interval.first_use_time = std::numeric_limits::max(); - use_interval.last_use_time = -1; - for (const HloUse& use : value->GetUses()) { - auto it = instruction_schedule.find(use.instruction); - if (it == instruction_schedule.end()) { - continue; - } - use_interval.first_use_time = - std::min(use_interval.first_use_time, it->second); - int64_t last_use_time = it->second; - if (view_color.has_value() && use.instruction->shape().has_layout() && - use.instruction->shape().layout().memory_space() == *view_color) { - last_use_time = - std::max(last_use_time, - ViewExtendedTransitiveUseTime(use.instruction, *view_color, - instruction_schedule)); - } - use_interval.last_use_time = - std::max(use_interval.last_use_time, last_use_time); - } - CHECK_NE(use_interval.last_use_time, -1); - } - return value_to_use_intervals; -} - -absl::flat_hash_set GetParameterInstructionsAliasedToOutput( - const HloInputOutputAliasConfig& alias_config, - const HloInstruction* root_instruction) { - absl::flat_hash_set aliased_parameter_positions; - alias_config.ForEachAlias([&](const ShapeIndex& output_index, - const HloInputOutputAliasConfig::Alias& alias) { - HloInstruction* parameter_instruction = - root_instruction->parent()->parameter_instruction( - alias.parameter_number); - aliased_parameter_positions.insert( - {parameter_instruction, alias.parameter_index}); - }); - return aliased_parameter_positions; -} - -// Marks all RepackAllocationBlocks in the list as colocated, forming a circular -// linked list, representing colocated allocations. -void MarkRepackAllocationBlocksColocated( - std::vector& colocations) { - if (colocations.empty()) { - return; - } - for (size_t i = 0; i < colocations.size() - 1; ++i) { - colocations[i]->next_colocated = colocations[i + 1]; - } - colocations.back()->next_colocated = colocations.front(); -} - -void PopulateExistingBlockPrefetchedValues( - const Options& options, const HloAliasAnalysis& alias_analysis, - std::vector& block_prefetched_values, - absl::flat_hash_map& - prefetch_done_value_to_original_value, - absl::flat_hash_map& - prefetch_done_value_to_prefetch_start_instruction) { - // Block prefetched values are prefetch done HloValues. Original values are - // the HLO values that are being prefetched. We maintain a map of prefetch - // done value to original value. We also maintain a map of prefetch done - // HloValue to the prefetch start instruction. - for (const auto& [position, custom_call_prefetch_details] : - options.hlo_position_to_custom_call_prefetch_details) { - const HloValue* original_value = - &alias_analysis.dataflow_analysis().GetUniqueValueAt( - position.instruction, position.index); - for (const auto& custom_call_prefetch_detail : - custom_call_prefetch_details) { - HloInstruction* prefetch_start = - custom_call_prefetch_detail.prefetch_start; - HloInstruction* prefetch_done = custom_call_prefetch_detail.prefetch_done; - const HloValue* prefetch_done_value = - &alias_analysis.dataflow_analysis().GetUniqueValueAt(prefetch_done, - {}); - prefetch_done_value_to_original_value[prefetch_done_value] = - original_value; - prefetch_done_value_to_prefetch_start_instruction[prefetch_done_value] = - prefetch_start; - block_prefetched_values.push_back(prefetch_done_value); - } - } -} - -} // namespace - -absl::Status MsaAlgorithm::AllocateAndScheduleExistingBlockPrefetches( - int64_t block_prefetching_starting_offset) { - if (options_.hlo_position_to_custom_call_prefetch_details.empty()) { - return absl::OkStatus(); - } - if (options_.reserved_bytes_for_block_prefetches <= 0) { - return absl::FailedPreconditionError( - "Reserved bytes for block prefetches is zero, we need memory to " - "schedule custom call block prefetches."); - } - // List of all block prefetched HloValues. - std::vector block_prefetched_values; - absl::flat_hash_map - prefetch_done_value_to_original_value; - absl::flat_hash_map - prefetch_done_value_to_prefetch_start_instruction; - - PopulateExistingBlockPrefetchedValues( - options_, alias_analysis_, block_prefetched_values, - prefetch_done_value_to_original_value, - prefetch_done_value_to_prefetch_start_instruction); - - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - - // Compute the live ranges for each block prefetched value. - absl::flat_hash_map value_to_use_intervals = - GetUseIntervals(block_prefetched_values, instruction_schedule, - options_.dus_view_color); - - // Erase all the values from block_prefetched_values that have been finalized. - block_prefetched_values.erase( - std::remove_if( - block_prefetched_values.begin(), block_prefetched_values.end(), - [&](const HloValue* value) { return IsValueFinalized(value); }), - block_prefetched_values.end()); - - // Sort block prefetched values in ascending order of first use time. - absl::c_sort(block_prefetched_values, - [&](const HloValue* a, const HloValue* b) { - return value_to_use_intervals.at(a).first_use_time < - value_to_use_intervals.at(b).first_use_time; - }); - - // All block prefetches should be placed within this limit. - int64_t block_prefetching_limit_bytes = - block_prefetching_starting_offset + - options_.reserved_bytes_for_block_prefetches; - CHECK_LE(block_prefetching_limit_bytes, options_.max_size_in_bytes); - VLOG(1) << "block prefetched values bytes limit: " - << block_prefetching_limit_bytes; - int64_t previous_start_time = -1; - int64_t max_in_flight_prefetches_allowed = - options_.max_outstanding_block_prefetches; - std::vector prefetch_end_times; - std::vector prefetch_done_schedule_before_times; - - absl::flat_hash_map value_to_pinned_allocation; - - // For each block prefetched value, we try to find a chunk within the block - // prefetching limit, ensuring FIFO ordering. After a suitable chunk is found - // for a block prefetched value, we: - // 1. Commit the chunk to the alternate memory. - // 2. Update the operands in alternate memory map. - // 3. Add the copy done and copy start values to the finalized values set. - // 4. Add a repack allocation block to the repack allocation blocks list. - // 5. If the prefetch done value is aliased with values other than the - // prefetch start value, allocate the aliased values. - // 6. Serve the uses of the original value from the pinned allocation in the - // default memory. - // 7. Clear the pending chunks after the loop. - for (const HloValue* prefetch_done_value : block_prefetched_values) { - UseInterval use_interval = value_to_use_intervals.at(prefetch_done_value); - int64_t first_use_time = use_interval.first_use_time; - int64_t last_use_time = use_interval.last_use_time; - auto it = prefetch_done_value_to_original_value.find(prefetch_done_value); - CHECK(it != prefetch_done_value_to_original_value.end()); - const HloValue* original_value = it->second; - int64_t definition_time = - instruction_schedule.at(original_value->defining_instruction()); - int64_t end_time = last_use_time; - int64_t buffer_size = buffer_intervals_.at(prefetch_done_value).size; - int64_t earliest_start_time_candidate = - std::max(definition_time, previous_start_time); - CHECK_LE(earliest_start_time_candidate, first_use_time); - - // Find the earliest start time for which a chunk can be allocated for the - // block prefetched value. - std::optional optional_start_time = EarliestBlockPrefetchStartTime( - previous_start_time, definition_time, first_use_time, end_time, - buffer_size, block_prefetching_limit_bytes, - max_in_flight_prefetches_allowed, prefetch_done_schedule_before_times, - prefetch_end_times); - - if (!optional_start_time.has_value()) { - // Custom call block prefetches must be allocated in the alternate memory. - return absl::FailedPreconditionError(absl::StrCat( - "Could not find a chunk for block prefetched value: ", - prefetch_done_value->defining_position().ToString(), " buffer size: ", - buffer_size, " within limit: ", block_prefetching_limit_bytes)); - } - - int64_t start_time = optional_start_time.value(); - MsaBufferInterval interval = MsaBufferInterval{/*buffer=*/original_value, - /*size=*/buffer_size, - /*start=*/start_time, - /*end=*/end_time, - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk chunk_candidate = FindChunkCandidate(interval); - - // The chunk candidate should always be within the block prefetched values - // limit, otherwise we would have returned earlier. - CHECK_LE(chunk_candidate.chunk_end(), block_prefetching_limit_bytes); - - // Add a pinned allocation in the default memory to serve as the prev - // allocation for the copy allocation or extend the existing pinned - // allocation. - Allocation* pinned_allocation; - auto pinned_allocation_it = value_to_pinned_allocation.find(original_value); - if (pinned_allocation_it == value_to_pinned_allocation.end()) { - allocations_->push_back(std::make_unique( - original_value->defining_position(), MemorySpace::kDefault, - kDummyChunk, definition_time, end_time)); - pinned_allocation = allocations_->back().get(); - value_to_pinned_allocation[original_value] = pinned_allocation; - } else { - pinned_allocation = pinned_allocation_it->second; - pinned_allocation->Extend(end_time); - } - - HloInstruction* async_mem_op_start = - prefetch_done_value_to_prefetch_start_instruction[prefetch_done_value]; - HloInstruction* asyn_mem_op_done = prefetch_done_value->instruction(); - - // Add an async slice copy for the block prefetched value value. - AddAsyncCopyOrOtherMemOp( - /*prev_allocation=*/*pinned_allocation, - /*memory_space=*/MemorySpace::kAlternate, - /*chunk=*/chunk_candidate, - /*exclusive_start_time=*/InclusiveToExclusiveStartTime(start_time), - /*end_time=*/end_time, - /*copy_done_schedule_before_time=*/first_use_time, - /*allocations=*/allocations_, - /*aliased_offset=*/nullptr, - /*resource=*/0.0, - /*cross_program_prefetch_index=*/std::nullopt, - /*sync_mem_op=*/nullptr, - /*async_mem_op_start*/ async_mem_op_start, - /*async_mem_op_done=*/asyn_mem_op_done); - - previous_start_time = start_time; - auto const sorted_position = std::lower_bound( - prefetch_end_times.begin(), prefetch_end_times.end(), end_time); - prefetch_end_times.insert(sorted_position, end_time); - prefetch_done_schedule_before_times.push_back(first_use_time); - - // 1. Commit the chunk to the alternate memory. - AddToPendingChunks(interval, chunk_candidate); - - for (const HloUse& use : prefetch_done_value->GetUses()) { - if (use.instruction->parent() == - prefetch_done_value->instruction()->parent()) { - allocations_->back()->AddUse(use); - // 2. Update the operands in alternate memory map. - AddOperandToAlternateMemoryMap(use.instruction, use.operand_number, - use.operand_index); - } - } - - // 3. Add the copy done and copy start values to the finalized values set. - FinalizeValue(prefetch_done_value); - const HloValue* copy_start_value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - async_mem_op_start, {0}); - FinalizeValue(copy_start_value); - - // 4. Add a repack allocation block to the repack allocation blocks list. - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - start_time, end_time, chunk_candidate.size, chunk_candidate.offset, - allocations_->back().get())); - repack_allocation_blocks_.back().next_colocated = - &(repack_allocation_blocks_.back()); - const HloBuffer& buffer = - alias_analysis_.GetBufferContainingValue(*prefetch_done_value); - - // 5. If the prefetch done value is aliased with values other than the - // prefetch start value, allocate the aliased values. - if (buffer.values().size() > 2) { - ColocateAndFinalizeValuesAliasedToExistingBlockPrefetches( - prefetch_done_value, buffer, chunk_candidate, buffer_size, - &repack_allocation_blocks_.back(), instruction_schedule, - value_to_use_intervals, - prefetch_done_value_to_prefetch_start_instruction, - prefetch_end_times); - } - } - - // 6. Serve the uses of the original value from the pinned allocation in the - // default memory. - for (auto [_, original_value] : prefetch_done_value_to_original_value) { - // Finalize all values aliased to the original value. - // Note: We do not need to add pinned allocations for the aliased values, - // just finalizing them is sufficient to ensure that they will be served - // from default memory. - const HloBuffer& buffer = - alias_analysis_.GetBufferContainingValue(*original_value); - for (const HloValue* aliased_value : buffer.values()) { - if (IsValueFinalized(aliased_value)) { - continue; - } - // If a pinned allocation already exists for the aliased value, add the - // uses of the original value to the pinned allocation. - auto it = value_to_pinned_allocation.find(aliased_value); - if (it != value_to_pinned_allocation.end()) { - Allocation* allocation = it->second; - for (const HloUse& use : original_value->GetUses()) { - if (use.instruction->parent() == - original_value->instruction()->parent()) { - allocation->AddUse(use); - } - } - } - FinalizeValue(aliased_value); - } - } - - // 7. Clear the pending chunks. - ClearPendingChunks(); - return absl::OkStatus(); -} - -void MsaAlgorithm::ColocateAndFinalizeValuesAliasedToExistingBlockPrefetches( - const HloValue* prefetch_done_value, const HloBuffer& buffer, - const Chunk& chunk_candidate, int64_t buffer_size, - AllocationBlock* first_colocated_repack_allocation, - const absl::flat_hash_map& - instruction_schedule, - const absl::flat_hash_map& - value_to_use_intervals, - absl::flat_hash_map& - prefetch_done_value_to_prefetch_start_instruction, - std::vector& prefetch_end_times) { - VLOG(1) << "HloBuffer for block prefetched value: " - << prefetch_done_value->ToShortString() - << " aliases with: " << (buffer.values().size() - 1) - << " other values"; - - std::vector colocated_values; - for (const HloValue* aliased_value : buffer.values()) { - colocated_values.push_back(aliased_value); - } - absl::c_sort(colocated_values, [&](const HloValue* a, const HloValue* b) { - return instruction_schedule.at(a->defining_instruction()) < - instruction_schedule.at(b->defining_instruction()); - }); - - // After sorting by schedule time, the first two values in the colocated - // values list should be the prefetch start and prefetch done values, followed - // by their uses which might be aliased. - CHECK_EQ( - colocated_values[0]->instruction(), - prefetch_done_value_to_prefetch_start_instruction[prefetch_done_value]); - CHECK_EQ(colocated_values[1], prefetch_done_value); - - int64_t prev_last_use_time = - value_to_use_intervals.at(prefetch_done_value).last_use_time; - - std::vector colocations; - colocations.push_back(first_colocated_repack_allocation); - - int64_t maybe_sliced_value_definition_time = - instruction_schedule.at(prefetch_done_value->defining_instruction()); - - // For each of the colocated values that follow the block prefetched value, - // extend the chunk candidate to the right and add a pinned allocation in - // the alternate memory. We start from index 2, since the first two values - // in the colocated values list are the prefetch start and prefetch done - // values. - for (int i = 2; i < colocated_values.size(); ++i) { - const HloValue* aliased_value = colocated_values[i]; - CHECK(!IsValueFinalized(aliased_value)); - int64_t aliased_value_definition_time = - instruction_schedule.at(aliased_value->defining_instruction()); - CHECK_LT(maybe_sliced_value_definition_time, aliased_value_definition_time); - // The last use time of the previous value in the colocated values list - // should be the definition time of the current value in the colocated - // values list. This is because only the last use of a value can be - // aliased. - CHECK_EQ(prev_last_use_time, aliased_value_definition_time); - int64_t aliased_value_last_use_time = std::numeric_limits::min(); - for (const HloUse& use : aliased_value->GetUses()) { - aliased_value_last_use_time = - std::max(aliased_value_last_use_time, - instruction_schedule.at(use.instruction)); - } - prev_last_use_time = aliased_value_last_use_time; - - MsaBufferInterval aliased_interval = MsaBufferInterval{ - /*buffer=*/aliased_value, - /*size=*/buffer_size, - /*start=*/aliased_value_definition_time + - 1, // We need to add 1 because a chunk is already reserved till - // the prev_last_use_time which is equal to the - // aliased_value_definition_time. - /*end=*/aliased_value_last_use_time, - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk aliased_chunk_candidate = FindChunkCandidate( - aliased_interval, /*preferred_offset=*/chunk_candidate.offset); - // The aliased chunk candidate should be the same as the chunk candidate, - // since they are colocated and aliased. We are in principle extending the - // same chunk candidate to the right and we should always be able to do - // that because we are processing the values from left to right and we - // have checked that the prefetched value is only aliased to the right. - CHECK_EQ(aliased_chunk_candidate, chunk_candidate); - allocations_->push_back(std::make_unique( - aliased_value->defining_position(), MemorySpace::kAlternate, - aliased_chunk_candidate, aliased_value_definition_time, - aliased_value_last_use_time)); - AddToPendingChunks(aliased_interval, aliased_chunk_candidate); - for (const HloUse& use : aliased_value->GetUses()) { - if (use.instruction->parent() == aliased_value->instruction()->parent()) { - allocations_->back()->AddUse(use); - AddOperandToAlternateMemoryMap(use.instruction, use.operand_number, - use.operand_index); - } - } - auto const sorted_position = - std::lower_bound(prefetch_end_times.begin(), prefetch_end_times.end(), - aliased_value_last_use_time); - prefetch_end_times.insert(sorted_position, aliased_value_last_use_time); - FinalizeValue(aliased_value); - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - aliased_value_definition_time, aliased_value_last_use_time, - aliased_chunk_candidate.size, aliased_chunk_candidate.offset, - allocations_->back().get())); - repack_allocation_blocks_.back().next_colocated = - &(repack_allocation_blocks_.back()); - colocations.push_back(&repack_allocation_blocks_.back()); - } - - // Mark repack allocation blocks as colocated. - MarkRepackAllocationBlocksColocated(colocations); -} - -bool MsaAlgorithm::IsBufferAliasedToProgramOutput( - const HloBuffer* buffer) const { - for (const HloPosition& position : buffer->ComputePositions()) { - if (position.instruction == - module_->entry_computation()->root_instruction()) { - return true; - } - } - return false; -} - -bool MsaAlgorithm::DoesValueHaveNonAsyncConversionCandidateUses( - const HloValue* value) const { - for (const HloUse& use : value->GetUses()) { - if (!IsAsyncConversionCandidate(use.instruction)) { - return true; - } - } - return false; -} - -namespace { - -bool IsNestedInsideConditional(HloInstruction* instruction) { - HloComputation* computation = instruction->parent(); - for (HloInstruction* caller : computation->caller_instructions()) { - if (caller->opcode() == HloOpcode::kConditional || - IsNestedInsideConditional(caller)) { - return true; - } - } - return false; -} - -const HloValue* GetEarliestDefiningValue( - const HloBuffer* buffer, - const absl::flat_hash_map& - instruction_schedule) { - const HloValue* first_value = nullptr; - int64_t earliest_definition_time = std::numeric_limits::max(); - for (const HloValue* value : buffer->values()) { - auto it = instruction_schedule.find(value->instruction()); - if (it == instruction_schedule.end()) { - continue; - } - int64_t def_time = it->second; - if (def_time < earliest_definition_time) { - earliest_definition_time = def_time; - first_value = value; - } - } - return first_value; -} - -} // namespace - -const MsaAlgorithm::BlockPrefetchContext -MsaAlgorithm::BuildBlockPrefetchingContext( - const absl::flat_hash_map& - instruction_schedule) const { - BlockPrefetchCandidateAnalysis initial_analysis = - InitializeBlockPrefetchCandidateAnalysis(instruction_schedule); - - FilterAsyncConversionCandidates(initial_analysis); - - FilterBlockPrefetchCandidatesByUseConstraints(initial_analysis, - instruction_schedule); - - return BuildBlockPrefetchingContextHelper(std::move(initial_analysis), - instruction_schedule); -} - -const HloBuffer* MsaAlgorithm::GetScheduledBuffer( - const HloPosition& position, - const absl::flat_hash_map& - instruction_schedule) const { - if (!instruction_schedule.contains(position.instruction)) { - VLOG(3) << "Not block prefetching buffer because instruction is not in the " - "instruction schedule: " - << position.instruction->ToString(); - return nullptr; - } - return &alias_analysis_.GetUniqueBufferAt(position.instruction, - position.index); -} - -MsaAlgorithm::BlockPrefetchCategory MsaAlgorithm::ClassifyBlockPrefetchBuffer( - const HloBuffer* buffer, const HloValue* first_defining_value) const { - if (IsBufferAliasedToProgramOutput(buffer)) { - VLOG(3) << "Not block prefetching buffer aliased to program output: " - << buffer->ToString(); - return BlockPrefetchCategory::kPinToDefaultMemory; - } - - if (!DoesValueHaveNonAsyncConversionCandidateUses(first_defining_value)) { - VLOG(3) << "Not block prefetching buffer because its first defining " - "value does not have any uses that benefit from a copy of the " - "HloValue in Vmem: " - << buffer->ToString(); - return BlockPrefetchCategory::kPinToDefaultMemory; - } - - HloInstruction* first_instruction = first_defining_value->instruction(); - - // Program inputs need no further checks! - if (first_instruction->opcode() == HloOpcode::kParameter && - first_instruction->parent() == module_->entry_computation()) { - VLOG(3) << "Trying to block prefetch a program input buffer: " - << buffer->ToString(); - return BlockPrefetchCategory::kBlockPrefetchOnly; - } - - if (!IsAsyncConversionCandidate(first_instruction)) { - VLOG(3) - << "Not block prefetching buffer because first value of the " - "buffer is not an async conversion candidate or a program input: " - << buffer->ToString(); - return BlockPrefetchCategory::kPinToDefaultMemory; - } - - // If the async conversion candidate instruction is nested inside a - // conditional and the buffer is also the output of the conditional, we do not - // block prefetch it. This is because it may impose infeasible constraints on - // the allocations in the other branches of the conditional. - // TODO(b/535296054): Allow block prefetching if the buffer is not the output - // of the conditional. - if (IsNestedInsideConditional(first_instruction)) { - VLOG(3) - << "Not block prefetching async conversion candidate buffer because " - "its first value is nested inside a conditional: " - << buffer->ToString(); - return BlockPrefetchCategory::kPinToDefaultMemory; - } - - VLOG(3) << "Considering block prefetching buffer for async conversion: " - << buffer->ToString(); - return BlockPrefetchCategory::kAsyncConversionCandidate; -} - -MsaAlgorithm::AsyncConversionMetadata -MsaAlgorithm::ComputeBlockPrefetchingAsyncConversionMetadata( - HloInstruction* first_instruction, - const absl::flat_hash_map& - instruction_schedule) const { - // By default, we assume that the source buffer for every async conversion - // candidate instruction is operand 0 at {} shape index. However, this can be - // overridden by options_.custom_fusion_block_prefetch_operand_index_fn. - int64_t operand_number = 0; - if (first_instruction->IsCustomFusion()) { - std::optional custom_fusion_index = - options_.custom_fusion_block_prefetch_operand_index_fn( - first_instruction); - if (custom_fusion_index.has_value()) { - operand_number = custom_fusion_index.value(); - } - } - - const HloBuffer& source_buffer = alias_analysis_.GetUniqueBufferAt( - first_instruction->mutable_operand(operand_number), {}); - HloPosition non_trivial_source_position = GetNonTrivialSourcePosition( - {first_instruction->mutable_operand(operand_number), {}}); - - return AsyncConversionMetadata{&source_buffer, non_trivial_source_position}; -} - -MsaAlgorithm::BlockPrefetchCandidateAnalysis -MsaAlgorithm::InitializeBlockPrefetchCandidateAnalysis( - const absl::flat_hash_map& - instruction_schedule) const { - BlockPrefetchCandidateAnalysis analysis; - for (const HloPosition& position : options_.block_prefetched_positions) { - const HloBuffer* buffer = - GetScheduledBuffer(position, instruction_schedule); - if (!buffer) { - continue; - } - - const HloValue* first_defining_value = - GetEarliestDefiningValue(buffer, instruction_schedule); - if (first_defining_value != nullptr) { - analysis.buffer_to_first_defining_value[buffer] = first_defining_value; - } - - BlockPrefetchCategory category = - ClassifyBlockPrefetchBuffer(buffer, first_defining_value); - analysis.buffer_categories[buffer] = category; - - if (category != BlockPrefetchCategory::kAsyncConversionCandidate) { - continue; - } - - HloInstruction* first_instruction = first_defining_value->instruction(); - AsyncConversionMetadata metadata = - ComputeBlockPrefetchingAsyncConversionMetadata(first_instruction, - instruction_schedule); - analysis.async_candidate_metadata[buffer] = metadata; - analysis.buffer_to_first_defining_value[metadata.source_buffer] = - GetEarliestDefiningValue(metadata.source_buffer, instruction_schedule); - } - return analysis; -} - -void MsaAlgorithm::FilterAsyncConversionCandidates( - BlockPrefetchCandidateAnalysis& initial_analysis) const { - for (const auto& [candidate, metadata] : - initial_analysis.async_candidate_metadata) { - BlockPrefetchCategory& candidate_category = - initial_analysis.buffer_categories.at(candidate); - if (candidate_category == BlockPrefetchCategory::kPinToDefaultMemory) { - continue; - } - - const HloBuffer* src_buffer = metadata.source_buffer; - - // TODO(b/537894642): Instead of pinning the candidate to default memory, - // should we have a new category, where we let the defining instruction - // remain synchronous and allow the result to be pinned to alternate memory - // later? - // - // We do not allow DMA from alternate memory to alternate memory. If both - // the destination and source buffers are async conversion candidates we can - // only allow one of them to be block prefetched to alternate memory. - // Currently, we choose to block prefetch the source buffer over the - // destination buffer. - auto src_category_it = initial_analysis.buffer_categories.find(src_buffer); - if (src_category_it != initial_analysis.buffer_categories.end() && - src_category_it->second == - BlockPrefetchCategory::kAsyncConversionCandidate) { - VLOG(3) << "Not block prefetching an async conversion candidate buffer " - "as its source buffer is also an async conversion candidate: " - << candidate->ToString(); - candidate_category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - - // TODO(b/537894642): Should we also check if the src buffer is not an - // async conversion candidate? - // - // If the source buffer is a regular buffer (non-block-prefetched buffer), - // no more checks are required. - if (src_category_it == initial_analysis.buffer_categories.end() || - src_category_it->second != BlockPrefetchCategory::kBlockPrefetchOnly) { - initial_analysis.buffer_categories[src_buffer] = - BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - - // A block prefetched buffer has its first allocation as a pinned allocation - // in the default memory, followed by a copy allocation in the alternate - // memory, followed by pinned or mirrored allocations in the alternate - // memory for other positions. - // Since we do not allow DMA from alternate memory to alternate memory, if - // the source buffer of an async conversion candidate is block prefetched, - // we check if async conversion candidate's non-trivial source position is - // the defining position of the source buffer. If so, we can serve it from - // the pinned allocation in default memory, ensuring a DMA from default - // memory to alternate memory. - if (metadata.non_trivial_source_position != - initial_analysis.buffer_to_first_defining_value.at(src_buffer) - ->defining_position()) { - VLOG(3) << "Not block prefetching an async conversion candidate buffer " - "because its source buffer is also block prefetched and " - "the non trivial source position is not the defining position " - "of the source buffer." - << candidate->ToString(); - candidate_category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - - // The candidate retains its async conversion candidacy. - VLOG(3) - << "Will try to block prefetch an async conversion candidate buffer: " - << candidate->ToString(); - } -} - -MsaAlgorithm::BlockPrefetchCandidateUseAnalysis -MsaAlgorithm::ComputeBlockPrefetchCandidateUseAnalysis( - const HloBuffer* buffer, - const absl::flat_hash_map& - instruction_schedule, - const absl::flat_hash_set& - async_conv_candidate_instructions) const { - BlockPrefetchCandidateUseAnalysis analysis; - for (const HloValue* value : buffer->values()) { - for (const HloUse& use : value->GetUses()) { - if (async_conv_candidate_instructions.contains(use.instruction)) { - continue; - } - auto it = instruction_schedule.find(use.instruction); - if (it == instruction_schedule.end()) { - continue; - } - // Skip trivial uses that are not the root instruction. - if (IsUseTrivialNonRoot(use)) { - continue; - } - if (!options_.is_use_allowed_in_alternate_mem_fn(use)) { - analysis.all_uses_allowed_in_alternate_memory = false; - break; - } - if (HloDataflowAnalysis::IsAsynchronousOperationStart( - use.instruction->opcode())) { - analysis.all_uses_are_synchronous = false; - break; - } - int64_t corrected_use_time = GetCorrectedUseTime(use); - if (corrected_use_time < analysis.first_use_time) { - analysis.first_use_time = corrected_use_time; - analysis.first_use_instruction = use.instruction; - } - } - if (!analysis.all_uses_allowed_in_alternate_memory) { - break; - } - } - return analysis; -} - -bool MsaAlgorithm::HasDataMovementPositionsInMiddle( - const HloBuffer* buffer, const HloInstruction* first_instruction) const { - for (const HloPosition& position : buffer->ComputePositions()) { - if (position.instruction == first_instruction || - IsTrivialPosition(position)) { - continue; - } - HloOpcode op_code = position.instruction->opcode(); - if (op_code == HloOpcode::kCopy || op_code == HloOpcode::kSlice || - op_code == HloOpcode::kDynamicSlice || - op_code == HloOpcode::kCopyStart) { - return true; - } - } - return false; -} - -void MsaAlgorithm::FilterBlockPrefetchCandidatesByUseConstraints( - BlockPrefetchCandidateAnalysis& initial_analysis, - const absl::flat_hash_map& - instruction_schedule) const { - absl::flat_hash_set async_conv_candidate_instructions; - for (const auto& [buffer, category] : initial_analysis.buffer_categories) { - if (category == BlockPrefetchCategory::kAsyncConversionCandidate) { - async_conv_candidate_instructions.insert( - initial_analysis.buffer_to_first_defining_value.at(buffer) - ->instruction()); - } - } - - for (auto& [buffer, category] : initial_analysis.buffer_categories) { - if (category == BlockPrefetchCategory::kPinToDefaultMemory) { - continue; - } - - BlockPrefetchCandidateUseAnalysis use_analysis = - ComputeBlockPrefetchCandidateUseAnalysis( - buffer, instruction_schedule, async_conv_candidate_instructions); - if (use_analysis.first_use_time == std::numeric_limits::max()) { - VLOG(3) << "Not block prefetching a buffer because it has no uses: " - << buffer->ToString(); - category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - if (!use_analysis.all_uses_allowed_in_alternate_memory) { - VLOG(3) - << "Not block prefetching a buffer because not all uses are allowed " - "in alternate memory: " - << buffer->ToString(); - category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - if (!use_analysis.all_uses_are_synchronous) { - VLOG(3) << "Not block prefetching a buffer because not all uses are " - "synchronous: " - << buffer->ToString(); - category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - - const HloInstruction* first_instruction = - initial_analysis.buffer_to_first_defining_value.at(buffer) - ->instruction(); - if (HasDataMovementPositionsInMiddle(buffer, first_instruction)) { - VLOG(3) << "Not block prefetching a buffer because it has data movement " - "positions in the middle: " - << buffer->ToString(); - category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - - if (ShapeUtil::IsEffectiveScalar(first_instruction->shape())) { - VLOG(3) << "Not block prefetching a buffer because it is an effective " - "scalar: " - << buffer->ToString(); - category = BlockPrefetchCategory::kPinToDefaultMemory; - continue; - } - } -} - -MsaAlgorithm::BlockPrefetchContext -MsaAlgorithm::BuildBlockPrefetchingContextHelper( - BlockPrefetchCandidateAnalysis initial_analysis, - const absl::flat_hash_map& - instruction_schedule) const { - absl::flat_hash_set async_conv_candidate_instructions; - for (const auto& [buffer, category] : initial_analysis.buffer_categories) { - if (category == BlockPrefetchCategory::kAsyncConversionCandidate) { - async_conv_candidate_instructions.insert( - initial_analysis.buffer_to_first_defining_value.at(buffer) - ->instruction()); - } - } - - std::vector block_prefetched_buffers_list; - absl::flat_hash_map buffer_to_first_use_time; - - for (const auto& [buffer, category] : initial_analysis.buffer_categories) { - if (category != BlockPrefetchCategory::kPinToDefaultMemory) { - block_prefetched_buffers_list.push_back(buffer); - BlockPrefetchCandidateUseAnalysis use_analysis = - ComputeBlockPrefetchCandidateUseAnalysis( - buffer, instruction_schedule, async_conv_candidate_instructions); - VLOG(3) << "First use time for buffer " << buffer->ToString() << " is " - << use_analysis.first_use_time << " at instruction " - << use_analysis.first_use_instruction->name(); - buffer_to_first_use_time[buffer] = use_analysis.first_use_time; - } - } - - absl::c_sort(block_prefetched_buffers_list, [&](const HloBuffer* a, - const HloBuffer* b) { - return std::forward_as_tuple(buffer_to_first_use_time.at(a), a->id()) < - std::forward_as_tuple(buffer_to_first_use_time.at(b), b->id()); - }); - - return BlockPrefetchContext{std::move(initial_analysis), - std::move(block_prefetched_buffers_list), - std::move(buffer_to_first_use_time), - std::move(async_conv_candidate_instructions)}; -} - -Allocation* MsaAlgorithm::CreatePinnedAllocationInDefaultMemory( - HloPosition position, int64_t first_use_time, - const absl::flat_hash_map& - instruction_schedule, - BlockPrefetchSchedulingState& scheduling_state, - AllocationSequence& allocations) { - auto [pinned_allocation_it, inserted] = - scheduling_state.buffer_defining_source_position_to_pinned_allocation - .try_emplace(position, nullptr); - Allocation* pinned_allocation; - if (inserted) { - allocations.push_back(std::make_unique( - position, MemorySpace::kDefault, kDummyChunk, - instruction_schedule.at(position.instruction), first_use_time)); - pinned_allocation = allocations.back().get(); - pinned_allocation_it->second = pinned_allocation; - } else { - pinned_allocation = pinned_allocation_it->second; - pinned_allocation->Extend(first_use_time); - } - return pinned_allocation; -} - -void MsaAlgorithm::CreatePinnedAllocationsInAltMemoryForPositions( - absl::Span positions, const Chunk& chunk_candidate, - const absl::flat_hash_map& position_to_live_range, - const absl::flat_hash_map>& - position_to_uses, - const absl::flat_hash_set& - async_conv_candidate_instructions, - AllocationSequence& allocations, std::vector& colocations, - std::optional last_position_end_time) { - for (size_t idx = 0; idx < positions.size(); idx++) { - HloPosition non_trivial_position = positions[idx]; - LiveRange live_range = position_to_live_range.at(non_trivial_position); - int64_t start_time = live_range.start_time; - int64_t end_time = live_range.end_time; - if (idx < positions.size() - 1) { - end_time = position_to_live_range.at(positions[idx + 1]).start_time; - } else if (last_position_end_time.has_value()) { - end_time = last_position_end_time.value(); - } - allocations.emplace_back(std::make_unique( - non_trivial_position, MemorySpace::kAlternate, chunk_candidate, - start_time, end_time)); - Allocation* pinned_allocation = allocations.back().get(); - repack_allocation_blocks_.emplace_back( - MakeRepackAllocationBlock(start_time, end_time, chunk_candidate.size, - chunk_candidate.offset, pinned_allocation)); - colocations.push_back(&repack_allocation_blocks_.back()); - for (const HloUse* use : - FindOrDefault(position_to_uses, non_trivial_position, {})) { - CHECK(!async_conv_candidate_instructions.contains(use->instruction)); - CHECK_EQ(non_trivial_position, - GetNonTrivialSourcePosition( - {use->instruction->mutable_operand(use->operand_number), - use->operand_index})); - pinned_allocation->AddUse(*use); - AddOperandToAlternateMemoryMap(use->instruction, use->operand_number, - use->operand_index); - if (IsAsyncConversionCandidate(use->instruction)) { - failed_async_conversions_[use->instruction] = - AsyncConversionResult::kSourceBufferInAlternateMemory; - } - } - } -} - -absl::Status MsaAlgorithm::CreateMirroredAllocationsInAlternateMemory( - absl::Span positions, - const LiveRangeCalculator& calculator, - const absl::flat_hash_map& - instruction_schedule, - const absl::flat_hash_map>& - position_to_uses, - const absl::flat_hash_set& - async_conv_candidate_instructions, - AllocationSequence& allocations) { - for (const HloPosition& non_trivial_position : positions) { - HloInstruction* position_instruction = non_trivial_position.instruction; - int64_t start_time = - calculator.GetEarliestInstructionTime(position_instruction->parent()); - int64_t end_time = instruction_schedule.at( - position_instruction->parent()->root_instruction()); - Allocation* parent_allocation = nullptr; - int64_t instruction_time = instruction_schedule.at(position_instruction); - for (std::unique_ptr& allocation : allocations) { - if (allocation->is_in_default_mem()) { - continue; - } - if (allocation->is_mirrored_allocation()) { - continue; - } - if (allocation->start_time() <= instruction_time && - instruction_time <= allocation->end_time()) { - parent_allocation = allocation.get(); - break; - } - } - TF_RET_CHECK(parent_allocation != nullptr) - << "Parent allocation not found for mirrored allocation at: " - << non_trivial_position.ToString(); - allocations.emplace_back(std::make_unique( - non_trivial_position, *parent_allocation, start_time, end_time)); - Allocation* mirrored_allocation = allocations.back().get(); - for (const HloUse* use : - FindOrDefault(position_to_uses, non_trivial_position, {})) { - CHECK(!async_conv_candidate_instructions.contains(use->instruction)); - CHECK_EQ(non_trivial_position, - GetNonTrivialSourcePosition( - {use->instruction->mutable_operand(use->operand_number), - use->operand_index})); - mirrored_allocation->AddUse(*use); - AddOperandToAlternateMemoryMap(use->instruction, use->operand_number, - use->operand_index); - if (IsAsyncConversionCandidate(use->instruction)) { - failed_async_conversions_[use->instruction] = - AsyncConversionResult::kSourceBufferInAlternateMemory; - } - } - } - return absl::OkStatus(); -} - -// TODO(b/537894642): This method should be decomposed into smaller methods. -// -// Block prefetching a buffer is performed in the following steps: -// 1. Compute preliminaries. -// 2. Find the earliest time we can start prefetching the buffer. -// 3. Find an alternate memory chunk for the buffer. -// 4. Add a pinned allocation in the default memory to serve as the previous -// allocation for the (initial) copy allocation, or extend the existing -// pinned allocation. -// 5. For the first buffer position create a copy allocation. -// 6. Advance the global block prefetching scheduler state -// 7. Create repack allocation blocks for the allocations. -// 8. Complete the allocation sequence across all remaining positions in the -// buffer. -// 9. Finalize the block prefetch. -absl::StatusOr MsaAlgorithm::BlockPrefetchBuffer( - const HloBuffer* buffer, const HloValue* first_defining_value, - bool is_async_conversion_candidate, - std::optional async_non_trivial_source_position, - int64_t first_use_time, - const absl::flat_hash_map& - instruction_schedule, - int64_t block_prefetching_limit_bytes, - int64_t max_in_flight_prefetches_allowed, - const absl::flat_hash_set& - async_conv_candidate_instructions, - BlockPrefetchSchedulingState& scheduling_state) { - // 1. Compute preliminaries. - - // Compute live ranges for the buffer. - LiveRangeCalculator calculator(*buffer, instruction_schedule); - auto [buffer_live_ranges, has_overlap] = - calculator.CalculateBufferLiveRange(); - if (has_overlap) { - LOG(WARNING) - << "Not block prefetching a buffer because it has overlapping live " - "ranges: " - << buffer->ToString(); - return Decision::Forbid("Buffer has overlapping live ranges"); - } - const absl::flat_hash_map& position_to_live_range = - calculator.non_trivial_non_forwarded_position_to_live_range(); - const absl::flat_hash_map>& - position_to_uses = - calculator.non_trivial_non_forwarded_position_to_uses(); - - // Determine which positions need real allocations, which need mirrored - // allocations, and which can be ignored (because they are trivial). - std::vector non_trivial_positions_for_real_allocations; - std::vector non_trivial_positions_for_mirrored_allocations; - for (const HloPosition& position : buffer->ComputePositions()) { - if (IsTrivialPosition(position)) { - CHECK(!position_to_live_range.contains(position)); - continue; - } - if (position_to_live_range.contains(position)) { - non_trivial_positions_for_real_allocations.push_back(position); - } else { - non_trivial_positions_for_mirrored_allocations.push_back(position); - } - } - - // Sort the real positions chronologically by start time. - // In a well-formed buffer, the earliest position (front) corresponds to the - absl::c_stable_sort(non_trivial_positions_for_real_allocations, - [&](const HloPosition& a, const HloPosition& b) { - return position_to_live_range.at(a).start_time < - position_to_live_range.at(b).start_time; - }); - - int64_t live_range_start_time = buffer_live_ranges.front().start_time; - int64_t live_range_end_time = buffer_live_ranges.back().end_time; - int64_t data_readiness_time = live_range_start_time; - HloPosition buffer_defining_position = - first_defining_value->defining_position(); - CHECK(buffer_defining_position == - non_trivial_positions_for_real_allocations.front()); - - // If this buffer is an asynchronous conversion candidate, the async transfer - // cannot start until the source operand is computed. - if (is_async_conversion_candidate) { - VLOG(3) << "Block prefetched buffer: " << buffer->ToString() - << " is an async conversion candidate."; - std::optional latest_source_operand_time = - GetLatestSourceOperandScheduleTime(first_defining_value->instruction(), - instruction_schedule); - if (latest_source_operand_time.has_value()) { - data_readiness_time = - std::min(data_readiness_time, latest_source_operand_time.value()); - } - } - - // 2. Find the earliest time we can start prefetching the buffer. - int64_t buffer_size = buffer_intervals_.at(first_defining_value).size; - VLOG(3) << "Finding a chunk for block prefetched buffer: " - << buffer->ToString() << " buffer size: " << buffer_size - << " within limit: " << block_prefetching_limit_bytes; - std::optional earliest_block_prefetching_start_time = - EarliestBlockPrefetchStartTime( - scheduling_state.previous_start_time, data_readiness_time, - first_use_time, live_range_end_time, buffer_size, - block_prefetching_limit_bytes, max_in_flight_prefetches_allowed, - scheduling_state.prefetch_done_schedule_before_times, - scheduling_state.buffer_end_times); - if (!earliest_block_prefetching_start_time.has_value()) { - LOG(WARNING) << "Could not find a chunk for block prefetched buffer: " - << buffer->ToString() << " buffer size: " << buffer_size - << " within limit: " << block_prefetching_limit_bytes; - return Decision::Forbid("Could not find a chunk within limit"); - } - - // 3. Find an alternate memory chunk for the buffer. - int64_t allocation_start_time = earliest_block_prefetching_start_time.value(); - MsaBufferInterval interval = - MsaBufferInterval{/*buffer=*/nullptr, - /*size=*/buffer_size, - /*start=*/allocation_start_time, - /*end=*/live_range_end_time, - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk chunk_candidate = FindChunkCandidate(interval); - // We should always find a chunk candidate within the block prefetching - // bytes limit, because we already checked that when we found the earliest - // start time. - CHECK_LE(chunk_candidate.chunk_end(), block_prefetching_limit_bytes); - AddToPendingChunks(interval, chunk_candidate); - - // 4. Add a pinned allocation in the default memory to serve as the previous - // allocation for the copy allocation, or extend the existing pinned - // allocation. - AllocationSequence allocations; - std::vector colocations; - HloPosition non_trivial_source_position = buffer_defining_position; - HloInstruction* sync_mem_op = nullptr; - if (is_async_conversion_candidate) { - CHECK(async_non_trivial_source_position.has_value()); - non_trivial_source_position = *async_non_trivial_source_position; - sync_mem_op = first_defining_value->instruction(); - VLOG(3) << "Async conversion candidate block prefetched buffer: " - << buffer->ToString() << "\n non_trivial_source_position: " - << non_trivial_source_position.instruction->name() - << " sync_mem_op: " << sync_mem_op->name(); - } - Allocation* pinned_allocation = CreatePinnedAllocationInDefaultMemory( - non_trivial_source_position, first_use_time, instruction_schedule, - scheduling_state, allocations); - - // Determine when the initial asynchronous copy allocation terminates: - int64_t copy_allocation_end_time = - position_to_live_range - .at(non_trivial_positions_for_real_allocations.front()) - .end_time; - if (non_trivial_positions_for_real_allocations.size() > 1) { - copy_allocation_end_time = - position_to_live_range.at(non_trivial_positions_for_real_allocations[1]) - .start_time; - } - - // For asynchronous conversion candidates, identify the source operand index. - int64_t source_operand_index = 0; - if (sync_mem_op && sync_mem_op->IsCustomFusion()) { - auto custom_fusion_index = - options_.custom_fusion_block_prefetch_operand_index_fn(sync_mem_op); - CHECK(custom_fusion_index.has_value()); - source_operand_index = *custom_fusion_index; - } - - // 5. For the first buffer position create a copy allocation. - AddAsyncCopyOrOtherMemOp( - /*prev_allocation=*/*pinned_allocation, - /*memory_space=*/MemorySpace::kAlternate, - /*chunk=*/chunk_candidate, - /*exclusive_start_time=*/ - InclusiveToExclusiveStartTime(allocation_start_time), - /*end_time=*/copy_allocation_end_time, - /*copy_done_schedule_before_time=*/first_use_time, - /*allocations=*/&allocations, - /*aliased_offset=*/nullptr, - /*resource=*/0.0, - /*cross_program_prefetch_index=*/std::nullopt, - /*sync_mem_op=*/sync_mem_op, - /*async_mem_op_start=*/nullptr, - /*async_mem_op_done=*/nullptr, - /*source_operand_index=*/source_operand_index); - - // 6. Advance the global block prefetching scheduler state: record this - // prefetch's start and completion deadlines, and maintain buffer_end_times in - // sorted order so future block prefetch interval calculations can track - // active concurrency and in-flight byte limits. - scheduling_state.previous_start_time = allocation_start_time; - const auto sorted_position = std::lower_bound( - scheduling_state.buffer_end_times.begin(), - scheduling_state.buffer_end_times.end(), live_range_end_time); - scheduling_state.buffer_end_times.insert(sorted_position, - live_range_end_time); - scheduling_state.prefetch_done_schedule_before_times.push_back( - first_use_time); - - // 7. Create repack allocation blocks for the allocations. - Allocation* copy_allocation = allocations.back().get(); - repack_allocation_blocks_.emplace_back(MakeRepackAllocationBlock( - allocation_start_time, copy_allocation_end_time, chunk_candidate.size, - chunk_candidate.offset, copy_allocation)); - colocations.push_back(&repack_allocation_blocks_.back()); - - // Associate uses of the first position to the copy allocation. - for (const HloUse* use : position_to_uses.at(buffer_defining_position)) { - if (!async_conv_candidate_instructions.contains(use->instruction)) { - CHECK_EQ(buffer_defining_position, - GetNonTrivialSourcePosition( - {use->instruction->mutable_operand(use->operand_number), - use->operand_index})); - copy_allocation->AddUse(*use); - AddOperandToAlternateMemoryMap(use->instruction, use->operand_number, - use->operand_index); - if (IsAsyncConversionCandidate(use->instruction)) { - failed_async_conversions_[use->instruction] = - AsyncConversionResult::kSourceBufferInAlternateMemory; - } - } - } - - // 8. Complete the allocation sequence across all remaining positions in the - // buffer: - // - Create PinnedAllocations in alternate memory for all subsequent real - // positions (ensuring the buffer stays resident across non-trivial uses). - // - Create MirroredAllocations for aliased/phi positions to properly link - // their uses without allocating duplicate physical chunks. - CreatePinnedAllocationsInAltMemoryForPositions( - absl::MakeConstSpan(non_trivial_positions_for_real_allocations) - .subspan(1), - chunk_candidate, position_to_live_range, position_to_uses, - async_conv_candidate_instructions, allocations, colocations); - - ABSL_RETURN_IF_ERROR(CreateMirroredAllocationsInAlternateMemory( - non_trivial_positions_for_mirrored_allocations, calculator, - instruction_schedule, position_to_uses, async_conv_candidate_instructions, - allocations)); - - // 9. Finalize the block prefetch. - ClearPendingChunks(); - for (const HloValue* value : buffer->values()) { - FinalizeValue(value); - } - MarkRepackAllocationBlocksColocated(colocations); - for (auto& allocation : allocations) { - allocations_->push_back(std::move(allocation)); - } - - return Decision::Allow(); -} - -absl::Status MsaAlgorithm::CreateNewBlockPrefetches( - int64_t block_prefetching_starting_offset) { - if (!options_.hlo_position_to_custom_call_prefetch_details.empty() || - options_.reserved_bytes_for_block_prefetches <= 0 || - options_.block_prefetched_positions.empty()) { - return absl::OkStatus(); - } - VLOG(3) << "Trying to block prefetch " - << options_.block_prefetched_positions.size() << " positions."; - - const absl::flat_hash_map& - instruction_schedule = hlo_live_range_.instruction_schedule(); - const BlockPrefetchContext context = - BuildBlockPrefetchingContext(instruction_schedule); - - int64_t block_prefetching_limit_bytes = - block_prefetching_starting_offset + - options_.reserved_bytes_for_block_prefetches; - - if (block_prefetching_limit_bytes > options_.max_size_in_bytes) { - int64_t corrected_reserved_bytes = - options_.max_size_in_bytes - block_prefetching_starting_offset; - return absl::InvalidArgumentError(absl::StrCat( - "Block prefetched values bytes limit: ", block_prefetching_limit_bytes, - " is greater than options_.max_size_in_bytes ", - options_.max_size_in_bytes, ". Try setting it to ", - (corrected_reserved_bytes / 1024), " KiB.")); - } - VLOG(3) << "Block prefetched values bytes limit: " - << block_prefetching_limit_bytes; - - BlockPrefetchSchedulingState scheduling_state; - scheduling_state.previous_start_time = -1; - int64_t max_in_flight_prefetches_allowed = - options_.max_outstanding_block_prefetches; - - absl::flat_hash_set failed_prefetches; - for (const HloBuffer* buffer : context.block_prefetched_buffers_list) { - const HloValue* first_defining_value = - context.initial_analysis.buffer_to_first_defining_value.at(buffer); - bool is_async_conversion_candidate = - context.initial_analysis.buffer_categories.at(buffer) == - BlockPrefetchCategory::kAsyncConversionCandidate; - std::optional async_non_trivial_source_position; - if (is_async_conversion_candidate) { - async_non_trivial_source_position = - context.initial_analysis.async_candidate_metadata.at(buffer) - .non_trivial_source_position; - } - - ABSL_ASSIGN_OR_RETURN( - Decision decision, - BlockPrefetchBuffer( - buffer, first_defining_value, is_async_conversion_candidate, - async_non_trivial_source_position, - context.buffer_to_first_use_time.at(buffer), instruction_schedule, - block_prefetching_limit_bytes, max_in_flight_prefetches_allowed, - context.async_conv_candidate_instructions, scheduling_state)); - if (decision.IsForbidden()) { - failed_prefetches.insert(buffer); - } - } - - // NOLINTNEXTLINE - for (const auto& [buffer, category] : - context.initial_analysis.buffer_categories) { - if (category == BlockPrefetchCategory::kPinToDefaultMemory || - failed_prefetches.contains(buffer)) { - for (const HloValue* value : buffer->values()) { - FinalizeValue(value); - } - } - } - - return absl::OkStatus(); -} - -int64_t MsaAlgorithm::ReserveAlternateMemoryForScopedMemoryAllocations() { - int64_t max_scoped_memory_size = MaxScopedMemorySize(); - int program_start_time = 0; - int program_end_time = - hlo_live_range_.flattened_instruction_sequence().instructions().size() - - 1; - MsaBufferInterval interval; - interval.buffer = nullptr; - interval.size = max_scoped_memory_size; - interval.start = program_start_time; - interval.end = program_end_time; - interval.need_allocation = true; - HeapSimulator::Chunk chunk_candidate = - FindChunkCandidate(interval, /*preferred_offset=*/0); - CHECK_EQ(chunk_candidate.offset, 0); - CHECK_EQ(chunk_candidate.size, max_scoped_memory_size); - AddToPendingChunks(interval, chunk_candidate); - ClearPendingChunks(); - return max_scoped_memory_size; -} - -void MsaAlgorithm::FreeAlternateMemoryForScopedMemoryAllocations( - int64_t max_scoped_memory_size) { - int program_start_time = 0; - int program_end_time = - hlo_live_range_.flattened_instruction_sequence().instructions().size() - - 1; - HeapSimulator::Chunk chunk = HeapSimulator::Chunk::FromOffsetSize( - /*offset=*/0, /*size=*/max_scoped_memory_size); - MsaBufferInterval interval; - interval.buffer = nullptr; - interval.size = max_scoped_memory_size; - interval.start = program_start_time; - interval.end = program_end_time; - interval.need_allocation = true; - UncommitChunkAndUpdatePeakMemory(interval, chunk); -} - -void MsaAlgorithm::CalculateMemoryPressure( - const std::vector& sorted_buffer_intervals) { - // Calculate the memory pressure for the buffers that can be assigned in - // alternate memory. - memory_pressure_ = 0; - VLOG(5) << [&]() { - std::string s("Sorted MsaBufferInterval order."); - if (options_.buffer_interval_comparator) { - absl::StrAppend( - &s, " Pre-autotuning sort criteria: ", - options_.buffer_interval_comparator->DescribeComparisonCriteria()); - } - return s; - }(); - for (auto& interval : sorted_buffer_intervals) { - if (!interval.need_allocation || - !MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory( - interval, options_.alternate_memory_space) || - interval.size > available_heap_size()) { - continue; - } - VLOG(5) << [&]() { - std::string s("SortedBufferInterval."); - if (options_.buffer_interval_comparator) { - absl::StrAppend( - &s, " Criteria: ", - options_.buffer_interval_comparator->CriteriaToString(interval)); - } - absl::StrAppend(&s, " Buffer: ", interval.buffer->ToShortString()); - return s; - }(); - memory_pressure_ += interval.size; - } - VLOG(1) << "Memory pressure = " << memory_pressure_; -} - -void MsaAlgorithm::FinalizeValue(const HloValue* value) { - VLOG(2) << "Finalizing value: " << value->ToShortString(); - finalized_values_.insert(value); -} -bool MsaAlgorithm::IsBlockPrefetchingEnabled() const { - return !options_.block_prefetched_positions.empty() || - !options_.hlo_position_to_custom_call_prefetch_details.empty(); -} - -std::vector -MsaAlgorithm::GetPostProcessedSortedBufferIntervals() { - std::vector sorted_buffer_intervals = - GetSortedBufferIntervals(); - - // Customize the sorting of buffer intervals if an autotuning config is - // provided. - memory_space_assignment::CustomizeSortedBufferInterval( - options_.autotuning_config, sorted_buffer_intervals); - - // Apply other sorting constraints after the autotuning config is applied. - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - - auto get_definition_time = [&](const HloValue* value) { - return instruction_schedule.find(value->defining_instruction())->second; - }; - - auto get_last_use_time = [&](const HloValue* value) { - int64_t last_use_time = std::numeric_limits::min(); - for (const HloUse& use : value->GetUses()) { - last_use_time = std::max(last_use_time, GetCorrectedUseTime(use)); - } - return last_use_time; - }; - - auto is_pre_colored = [&](const HloValue* value) { - return value->shape().has_layout() && - value->shape().layout().memory_space() == - options_.alternate_memory_space; - }; - - auto does_buffer_have_alt_mem_coloring = [&](const HloValue* value) { - const HloBuffer& buffer = - alias_analysis_.GetUniqueBufferAt(value->defining_position()); - return reserved_allocations_for_alt_mem_colorings_.contains(&buffer); - }; - - struct BufferInfo { - bool pre_colored = false; - bool alt_mem_coloring = false; - bool async_pipelined_while_alias = false; - int64_t definition_time = std::numeric_limits::min(); - int64_t last_use_time = std::numeric_limits::min(); - }; - - // Detect whether the module contains async pipelined while loops (marked with - // frontend attributes controlling while loop copy elimination) to enable - // precise alternate memory offset colocation. - has_async_pipelined_while_loops_ = HasAsyncPipelinedWhileLoops(*module_); - - // Precompute buffer info for each buffer in O(n) time instead of calling the - // above functions repeatedly in the sort function. - absl::flat_hash_map buffer_info; - buffer_info.reserve(sorted_buffer_intervals.size()); - for (MsaBufferInterval& interval : sorted_buffer_intervals) { - BufferInfo& info = buffer_info[interval.buffer]; - info.pre_colored = is_pre_colored(interval.buffer); - if (has_async_pipelined_while_loops_) { - info.async_pipelined_while_alias = - IsBufferAliasedToAsyncPipelinedWhileLoop(interval.buffer); - } - } - - bool is_block_prefetching_enabled = - options_.reserved_bytes_for_block_prefetches > 0; - if (is_block_prefetching_enabled) { - // All prefetches will happen as a part of block prefetching, regular MSA - // will not do any prefetching and run in a pin-only mode. We need to sort - // the buffers in the following order: - // 1. Pre-colored buffers first. - // - Within pre-colored buffers, sort by definition time and last use - // time in that order. - // 2. Rest of the buffers. - // - Within these buffers, sort by size, definition time and last use - // time in that order. - - for (const MsaBufferInterval& interval : sorted_buffer_intervals) { - BufferInfo& properties = buffer_info[interval.buffer]; - properties.definition_time = get_definition_time(interval.buffer); - properties.last_use_time = get_last_use_time(interval.buffer); - } - - // TODO(b/442852498): Move this custom sorting logic to a new - // BufferIntervalComparator class. - absl::c_stable_sort( - sorted_buffer_intervals, - [&](const MsaBufferInterval& a, const MsaBufferInterval& b) { - // At least one of the buffers is pre-colored. Sort by pre-colored - // buffers first, then by definition time and last use time in that - // order. - BufferInfo& a_properties = buffer_info.at(a.buffer); - BufferInfo& b_properties = buffer_info.at(b.buffer); - if (a_properties.pre_colored || b_properties.pre_colored) { - return std::forward_as_tuple(!a_properties.pre_colored, - a_properties.definition_time, - a_properties.last_use_time) < - std::forward_as_tuple(!b_properties.pre_colored, - b_properties.definition_time, - b_properties.last_use_time); - } - // Both buffers are not pre-colored. Sort by size, definition time and - // last use time in that order. - return std::forward_as_tuple(a.size, a_properties.definition_time, - a_properties.last_use_time) < - std::forward_as_tuple(b.size, b_properties.definition_time, - b_properties.last_use_time); - }); - - return sorted_buffer_intervals; - } - - // If test mode is enabled, return the buffer intervals in the order they were - // returned by the GetSortedBufferIntervals() function. - if (!options_.allocate_colored_buffers_early) { - return sorted_buffer_intervals; - } - - // Stable sort the buffers in the following order: - // 1. Pre-colored buffers first. - // 2. Then, buffers with alt mem coloring first. - // Note the use of > instead of < in the sort function. - for (const MsaBufferInterval& interval : sorted_buffer_intervals) { - buffer_info[interval.buffer].alt_mem_coloring = - does_buffer_have_alt_mem_coloring(interval.buffer); - } - absl::c_stable_sort(sorted_buffer_intervals, [&](const MsaBufferInterval& a, - const MsaBufferInterval& b) { - BufferInfo& a_properties = buffer_info.at(a.buffer); - BufferInfo& b_properties = buffer_info.at(b.buffer); - if (has_async_pipelined_while_loops_) { - int64_t a_size = a_properties.async_pipelined_while_alias ? a.size : 0; - int64_t b_size = b_properties.async_pipelined_while_alias ? b.size : 0; - return std::forward_as_tuple( - a_properties.pre_colored, a_properties.alt_mem_coloring, - a_properties.async_pipelined_while_alias, a_size) > - std::forward_as_tuple( - b_properties.pre_colored, b_properties.alt_mem_coloring, - b_properties.async_pipelined_while_alias, b_size); - } - return std::forward_as_tuple(a_properties.pre_colored, - a_properties.alt_mem_coloring) > - std::forward_as_tuple(b_properties.pre_colored, - b_properties.alt_mem_coloring); - }); - - return sorted_buffer_intervals; -} - -absl::flat_hash_set -MsaAlgorithm::GetScalarBuffersPinnedToAltMemory( - absl::Span sorted_buffer_intervals) const { - absl::flat_hash_set scalar_buffers_to_pin_in_alt_mem; - for (const MsaBufferInterval& interval : sorted_buffer_intervals) { - HloInstruction* instruction = interval.buffer->instruction(); - if (ShapeUtil::IsEffectiveScalar(interval.buffer->shape()) && - instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == - options_.alternate_memory_space) { - CHECK(!IsValueFinalized(interval.buffer)) - << "Scalar pinned to alternate memory is already finalized: " - << interval.buffer->ToShortString(); - const HloBuffer& scalar_buffer = alias_analysis_.GetUniqueBufferAt( - instruction, interval.buffer->index()); - scalar_buffers_to_pin_in_alt_mem.insert(&scalar_buffer); - } - } - return scalar_buffers_to_pin_in_alt_mem; -} - -absl::Status MsaAlgorithm::PinScalarBuffersInAlternateMemory( - absl::flat_hash_set& scalar_buffers_to_pin_in_alt_mem) { - int64_t program_end_time = static_cast( - hlo_live_range_.flattened_instruction_sequence().instructions().size() - - 1); - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - std::vector sorted_scalar_buffers( - scalar_buffers_to_pin_in_alt_mem.begin(), - scalar_buffers_to_pin_in_alt_mem.end()); - absl::c_sort( - sorted_scalar_buffers, - [](const HloBuffer* a, const HloBuffer* b) { return a->id() < b->id(); }); - for (const HloBuffer* buffer : sorted_scalar_buffers) { - ABSL_RETURN_IF_ERROR(PinScalarBufferInAlternateMemory(buffer, program_end_time, - instruction_schedule)); - } - return absl::OkStatus(); -} - -absl::Status MsaAlgorithm::PinScalarBufferInAlternateMemory( - const HloBuffer* buffer, int64_t program_end_time, - const absl::flat_hash_map& - instruction_schedule) { - // For each scalar buffer pinned in alternate memory: - // 1. Compute live ranges for the buffer. - // 2. Partition non-trivial positions into positions that get real - // (pinned) allocations and positions that get mirrored allocations. - // 3. Find a chunk candidate for the buffer, skipping constant buffers or - // buffers that would OOM the alternate memory. - // 4. Commit the chunk to the alternate memory. - // 5. Create pinned allocations in alternate memory for the real positions, - // extending the last position to the end of the program (note: extending - // to program end is a temporary hack). - // 6. Create mirrored allocations in alternate memory for the remaining - // positions. - // 7. Finalize the buffer values. - // 8. Colocate the repack allocation blocks. - // 9. Add allocations to the allocations list. - // 10. Clear the pending chunks. - - // Compute live ranges of all positions in the buffer. - LiveRangeCalculator calculator(*buffer, instruction_schedule); - auto [buffer_live_ranges, has_overlap] = - calculator.CalculateBufferLiveRange(); - // outer_positions_to_live_range contain the minimal set of non-trivial - // positions required to capture the entire live range of a buffer without - // overlaps (i.e. nested non-trivial positions are not present in it). - const absl::flat_hash_map& - outer_positions_to_live_range = - calculator.non_trivial_non_forwarded_position_to_live_range(); - const absl::flat_hash_map>& - outer_positions_to_uses = - calculator.non_trivial_non_forwarded_position_to_uses(); - - // Partition non-trivial positions into positions that get real (pinned) - // allocations and positions that get mirrored allocations. - std::vector non_trivial_positions_for_real_allocations; - std::vector non_trivial_positions_for_mirrored_allocations; - for (const HloPosition& position : buffer->ComputePositions()) { - if (IsTrivialPosition(position)) { - TF_RET_CHECK(!outer_positions_to_live_range.contains(position)) - << "Trivial position should not have live range: " - << position.ToString(); - continue; - } - - if (outer_positions_to_live_range.contains(position)) { - non_trivial_positions_for_real_allocations.push_back(position); - } else { - non_trivial_positions_for_mirrored_allocations.push_back(position); - } - } - absl::c_stable_sort(non_trivial_positions_for_real_allocations, - [&](const HloPosition& a, const HloPosition& b) { - return outer_positions_to_live_range.at(a).start_time < - outer_positions_to_live_range.at(b).start_time; - }); - - // Find a chunk candidate for the buffer, skipping constant buffers or - // buffers that would OOM the alternate memory. - int64_t buffer_start_time = buffer_live_ranges.front().start_time; - const HloValue* first_defining_value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - non_trivial_positions_for_real_allocations.front().instruction, - non_trivial_positions_for_real_allocations.front().index); - const MsaBufferInterval& buffer_interval = - FindOrDefault(buffer_intervals_, first_defining_value, {}); - // A buffer might not be found in buffer_intervals_ (e.g. tokens, constants - // etc). - if (buffer_interval.buffer == nullptr || - first_defining_value->defining_instruction()->opcode() == - HloOpcode::kConstant) { - VLOG(2) << "Skipping coloring for buffer: " << buffer->ToString() - << " because it is a constant buffer or not found in buffer " - "intervals."; - return absl::OkStatus(); - } - int64_t buffer_size = buffer_interval.size; - MsaBufferInterval interval = - MsaBufferInterval{/*buffer=*/first_defining_value, - /*size=*/buffer_size, - /*start=*/buffer_start_time, - /*end=*/program_end_time, - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk chunk_candidate = FindChunkCandidate(interval); - if (chunk_candidate.chunk_end() > options_.max_size_in_bytes) { - return absl::FailedPreconditionError(absl::StrCat( - "Overcoloring scalars causes VMEM OOM. Scalar instruction: ", - first_defining_value->defining_instruction()->ToShortString())); - } - - // Commit the chunk to the alternate memory. - AddToPendingChunks(interval, chunk_candidate); - AllocationSequence allocations; - std::vector colocations; - - // Create pinned allocations in alternate memory for the real positions. - // Unlike block prefetching, scalar buffers have no leading copy allocation, - // so all real positions (starting at index 0) get pinned allocations, and - // the last position's allocation is extended to the end of the program. - CreatePinnedAllocationsInAltMemoryForPositions( - non_trivial_positions_for_real_allocations, chunk_candidate, - outer_positions_to_live_range, outer_positions_to_uses, - /*async_conv_candidate_instructions=*/{}, allocations, colocations, - /*last_position_end_time=*/program_end_time); - - // Create mirrored allocations in alternate memory for the remaining - // positions. - ABSL_RETURN_IF_ERROR(CreateMirroredAllocationsInAlternateMemory( - non_trivial_positions_for_mirrored_allocations, calculator, - instruction_schedule, outer_positions_to_uses, - /*async_conv_candidate_instructions=*/{}, allocations)); - - // Finalize the buffer values, colocate the repack allocation blocks, and add - // the allocations to the allocations list. - ClearPendingChunks(); - for (const HloValue* value : buffer->values()) { - FinalizeValue(value); - } - MarkRepackAllocationBlocksColocated(colocations); - for (auto& allocation : allocations) { - allocations_->push_back(std::move(allocation)); - } - return absl::OkStatus(); -} - -absl::StatusOr> MsaAlgorithm::Finish() { - // Note: Memory Space Assignment creates a HeapSimulator and passes an - // MsaAlgorithm object to it. buffer_intervals_ is populated by calling the - // Alloc(), Free() and ShareWith() methods on the MsaAlgorithm object in - // HeapSimulator. - if (options_.autotuning_config.has_value()) { - CHECK_EQ((*options_.autotuning_config).size(), buffer_intervals_.size()); - } - VLOG(1) << "Slicing is " - << (options_.sliced_prefetch_options.max_slices() >= 2 ? "enabled" - : "disabled"); - - // Reserve alternate memory for scoped memory allocations before block - // prefetches are allocated strictly above the MaxScopedMemorySize(). - int64_t max_scoped_memory_size = - ReserveAlternateMemoryForScopedMemoryAllocations(); - - ABSL_RETURN_IF_ERROR( - AllocateAndScheduleExistingBlockPrefetches(max_scoped_memory_size)); - ABSL_RETURN_IF_ERROR(CreateNewBlockPrefetches(max_scoped_memory_size)); - - // Free the alternate memory reserved for scoped memory allocations before - // allocating the scoped memory allocations. - FreeAlternateMemoryForScopedMemoryAllocations(max_scoped_memory_size); - AllocateReservedScopedAllocations(); - ABSL_RETURN_IF_ERROR(ProcessColoredBuffers()); - - std::vector sorted_buffer_intervals = - GetPostProcessedSortedBufferIntervals(); - - CalculateMemoryPressure(sorted_buffer_intervals); - - CrossProgramPrefetches cross_program_prefetches = FindCrossProgramPrefetches( - alias_analysis_, alias_info_, hlo_live_range_, options_); - // Return error if cross program prefetch is disabled and user has requested - // cross program prefetch. - if (!options_.enable_cross_program_prefetch && - !cross_program_prefetches.prefetches.empty()) { - return absl::FailedPreconditionError( - "Cross program prefetch is disabled but user has requested cross " - "program prefetch."); - } - - // Return error if number of user requested cross program prefetches is - // greater than the maximum number of cross program prefetches allowed. - if (cross_program_prefetches.prefetches.size() > - options().max_cross_program_prefetches) { - return absl::FailedPreconditionError( - "Number of user requested cross program prefetches is greater than the " - "maximum number of cross program prefetches allowed."); - } - - // Allocate user requested cross program prefetches first. - for (auto& prefetch : cross_program_prefetches.prefetches) { - HloModule* module = prefetch.buffer->instruction()->GetModule(); - AllocateCrossProgramPrefetchBuffer(module, prefetch); - } - if (options_.enable_cross_program_prefetch) { - for (auto& prefetch_candidate : cross_program_prefetches.candidates) { - HloModule* module = prefetch_candidate.buffer->instruction()->GetModule(); - if (0 <= options().max_cross_program_prefetches && - options().max_cross_program_prefetches <= - module->CrossProgramPrefetches().size()) { - break; - } - AllocateCrossProgramPrefetchBuffer(module, prefetch_candidate); - } - } - - VLOG(1) << "Assigning buffers to alternate memory. Max heap size = " - << options_.max_size_in_bytes; - - // Process colored buffers before input and output required assignments are - // added to avoid adding conflicting required assignments. - AddInputAndOutputRequiredAssignments(); - - if (VLOG_IS_ON(3) || options_.dump_fn != nullptr) { - VLOG(3) << "Flattened instruction sequence:"; - const auto& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - absl::StrAppend(&instruction_schedule_str_, "time,instruction_name\n"); - for (int i = 0; i < instruction_sequence.size(); ++i) { - VLOG(3) << " " << i << ": " << instruction_sequence[i]->parent()->name() - << " " << instruction_sequence[i]->name(); - absl::StrAppend(&instruction_schedule_str_, i, ",", - instruction_sequence[i]->name(), "\n"); - } - } - - if (options_.memory_bound_loop_optimizer_options.enabled()) { - IdentifyAndOptimizeMemoryBoundLoops(); - } - - for (const auto& interval : sorted_buffer_intervals) { - auto colocated_intervals = GetSortedColocatedIntervals(interval); - if (AreIntervalsReservedInAlternateMemory(colocated_intervals)) { - // Increment the reserved part of alternate memory so that it is not - // available for other buffers. - reserved_in_bytes_ += options_.size_fn(*interval.buffer); - } - } - - absl::flat_hash_set scalar_buffers_to_pin_in_alt_mem = - GetScalarBuffersPinnedToAltMemory(sorted_buffer_intervals); - ABSL_RETURN_IF_ERROR( - PinScalarBuffersInAlternateMemory(scalar_buffers_to_pin_in_alt_mem)); - - VLOG(2) << "Total reserved bytes = " << reserved_in_bytes_; - for (MsaBufferInterval& interval : sorted_buffer_intervals) { - VLOG(3) << "Processing buffer: " << interval.buffer->ToString(); - if (IsValueFinalized(interval.buffer)) { - VLOG(3) << "Skip entrance interval" << interval.buffer->ToShortString() - << " because it is already processed."; - continue; - } - - JointAllocationProposal proposal = GetJointProposal(interval); - if (proposal.allocation_values.empty()) { - VLOG(3) << "No allocation values for these joint-processed values." - << interval.buffer->ToString(); - continue; - } - // Retry allocating this value with larger limits if allocation fails. - bool repacked = false; - for (int retry_number = 0; retry_number < options_.max_retries; - retry_number++) { - for (auto& colocated_intervals : proposal.colocated_intervals) { - AddRequiredAssignmentsForConditionalOutputsIfNecessary( - colocated_intervals); - } - options_.prefetch_interval_picker->SetRetryNumber(retry_number); - ABSL_ASSIGN_OR_RETURN( - AllocationResult result, - AllocateAllocationValues(absl::MakeSpan(proposal.allocation_values))); - VLOG(2) << "Allocation result = " << ResultToString(result); - VLOG(4) - << "Non-finalized allocations after processing allocation values:"; - for (int allocation_value_idx = 0; - allocation_value_idx < proposal.allocation_values.size(); - ++allocation_value_idx) { - auto& allocation_value = - proposal.allocation_values.at(allocation_value_idx); - VLOG(4) << " " << allocation_value.ToShortString(); - for (auto& allocation : *allocation_value.allocation_sequence()) { - VLOG(4) << " " << allocation->ToString(); - } - } - if (result_is(result, AllocationResult::kFailSyncDataMoveReplacement)) { - if (!options_.enable_sync_copy_replacement && - !options_.enable_sync_slice_replacement) { - return absl::FailedPreconditionError( - "Allocation result is " - "AllocationResult::kFailSyncCopyReplacement, but " - "no sync replacement is enabled."); - } - - UncommitPendingWork(absl::MakeSpan(proposal.allocation_values)); - proposal = GetJointProposal(interval); - if (proposal.allocation_values.empty()) { - VLOG(3) - << "No allocation values found in the updated joint-processed " - "values. Moving on to the next set of joint-processed values."; - break; - } - --retry_number; - - } else if (result_requires_uncommit(result)) { - UncommitPendingWork(absl::MakeSpan(proposal.allocation_values)); - VLOG(2) << "Couldn't allocate. Retry number " << retry_number; - if (retry_number > 0 && !sorted_async_conversion_candidates_.empty()) { - failed_async_conversions_[sorted_async_conversion_candidates_.at(0)] = - AsyncConversionResult::kFailedGaveUp; - VLOG(2) << "Giving the allocation another chance by dropping one " - "async conversion candidate."; - proposal = GetJointProposal(interval); - --retry_number; - } - } else if ((result_is(result, AllocationResult::kFailOutOfMemory) || - options_.repack_after_every_allocation) && - num_repacks_ < options_.max_repacks && !repacked && - !RepackAllocationsIncludeConvertedSyncMemOp() && - // Disable buffer repacking when async pipelined while loops - // are present, as repacking can alter assigned physical - // offsets and break required aliasing offset colocation across - // loop boundaries. - !has_async_pipelined_while_loops_) { - UncommitPendingWork(absl::MakeSpan(proposal.allocation_values)); - ++num_repacks_; - repacked = true; - if (!options_.repacker) { - return absl::FailedPreconditionError("Repacker not provided."); - } - - std::vector repack_allocation_blocks; - ExportAllocationsForRepacking(repack_allocation_blocks); - VLOG(2) << "Repacking."; - auto repack_status = - options_.repacker->Repack(absl::MakeSpan(repack_allocation_blocks)); - CHECK_EQ(repack_status.status(), absl::OkStatus()); - VLOG(2) << "Repack complete. Modified = " << *repack_status; - // For debug and testing purpose, also update allocations if - // repack_after_every_allocation is on. - if (*repack_status || options_.repack_after_every_allocation) { - ImportRepackedAllocations(); - --retry_number; - } - if (*repack_status) { - ++num_repacks_successful_; - } - } else { - // Check if any of the allocation sites are inefficient. If so, get rid - // of the pending allocation, require all of the inefficient sites in - // the default memory, and perform allocation again. - std::vector inefficient_sites = {}; - if (sorted_async_conversion_candidates_.empty()) { - std::vector proposed_inefficient_sites = {}; - proposed_inefficient_sites = - GetInefficientAllocationSites(proposal.allocation_values); - // If any of the initial inefficient sites are colored in alternate - // memory, we respect the coloring, regardless of whether the site is - // inefficient. - for (const HloPositionOrUse& site : proposed_inefficient_sites) { - if (std::holds_alternative(site) && - IsPositionColoredInAlternateMemory( - std::get(site))) { - continue; - } - if (std::holds_alternative(site) && - IsUseColoredInAlternateMemory(std::get(site))) { - continue; - } - inefficient_sites.push_back(site); - } - } - // We try to add default memory requirements to inefficient sites so - // the allocator avoids an inefficient alternate memory placement on - // the next pass. - // - If we do add new requirements, we do not count the retry attempt - // and try again. - // - If we don't add new requirements, we fall through to commit the - // allocations and stop retrying. - if (!inefficient_sites.empty() && - InefficientSiteRetryCanProgress(inefficient_sites)) { - UncommitPendingWork(absl::MakeSpan(proposal.allocation_values)); - for (const HloPositionOrUse& site : inefficient_sites) { - // To avoid a livelock situation, we commit the required assignments - // right away. Otherwise, reallocation can find alternate memory - // allocations at other sites, which can also be inefficient. - std::visit( - [this](const auto& site) { - VLOG(3) << "Inefficient site: " << site.ToString(); - AddRequiredAssignment( - site, MemorySpace::kDefault, - RequiredMemoryAssignment::Source::kInefficientSite, - /*offset=*/nullptr, - /*add_to_pending=*/false); - }, - site); - } - --retry_number; - continue; - } - - // The allocation succeeded. Commit it and stop retrying. - FinalizeAllocations(absl::MakeSpan(proposal.allocation_values)); - break; - } - } - // Keep track of the processed values to prevent double-processing in future - // joint-processed intervals. - for (auto& value : proposal.values) { - FinalizeValue(value); - } - } - - CHECK(VerifyOperandsInAlternateMemoryMap(allocations_, - operands_in_alternate_memory_map_)) - << "operands_in_alternate_memory_map_ is not consistent with the " - "finalizied allocations."; - - // Skip post-allocation repacking when async pipelined while loops are - // present to prevent altering established alternate memory colocation - // offsets. - if (options_.repack_after_every_allocation && - !has_async_pipelined_while_loops_) { - if (!options_.repacker) { - return absl::FailedPreconditionError("Repacker cannot be null."); - } - - if (RepackAllocationsIncludeConvertedSyncMemOp()) { - return absl::InternalError( - "Repacking is not supported yet when there are converted sync mem " - "ops."); - } - - std::vector repack_allocation_blocks; - ExportAllocationsForRepacking(repack_allocation_blocks); - VLOG(2) << "Final Repacking."; - auto repack_status = - options_.repacker->Repack(absl::MakeSpan(repack_allocation_blocks)); - if (!repack_status.ok()) { - return repack_status.status(); - } - VLOG(2) << "Final Repack complete. Modified = " << *repack_status; - } - - if (options_.dump_fn != nullptr || VLOG_IS_ON(3)) { - for (auto& allocation : *allocations_) { - // Only fill allocation_info_str_ if needed. - AppendAllocationInfoDebugString(*allocation, allocation_info_str_); - } - } - - VLOG(1) << "Repack summary: " << num_repacks_successful_ - << " succeeded out of " << num_repacks_; - - VLOG(3) << "Debug buffer info: "; - XLA_VLOG_LINES(3, buffer_info_str_); - VLOG(3) << "Debug allocation info: "; - XLA_VLOG_LINES(3, allocation_info_str_); - DumpDebugStringsIfEnabled(); - - if (VLOG_IS_ON(3)) { - VLOG(3) << "Sync copy replacement summary: "; - VLOG(3) << "\tnumber of successful async conversion: " - << successful_async_conversion_set_.size(); - VLOG(3) << "\tnumber of failed async conversion: " - << failed_async_conversions_.size(); - for (const HloInstruction* inst : successful_async_conversion_set_) { - VLOG(3) << "Successful copy replacement: " << inst->ToString(); - } - for (auto& failure : failed_async_conversions_) { - VLOG(3) << "Failed copy replacement: " << failure.first->ToString() - << ", reason: " << int(failure.second); - } - } - - // Run post allocation transformation and fix the allocation sequence if - // needed. - if (options_.post_allocation_transformation_fn) { - auto has_in_place_user = [this](HloInstruction* instr) { - for (HloInstruction* user : instr->users()) { - auto alias_pairs = alias_info_->GetInPlaceInputOutputPairs(user); - for (const auto& [operand_index, output_index] : alias_pairs) { - if (user->operand(operand_index.operand_number) == instr) { - return true; - } - } - } - return false; - }; - VLOG(3) << "Running post allocation transformation on module"; - for (HloComputation* comp : alias_analysis_.dataflow_analysis() - .module() - .MakeNonfusionComputations()) { - for (HloInstruction* instr : comp->MakeInstructionPostOrder()) { - // If the operand is in alternate memory, we don't run the - // post-allocation transformation. - auto operand_it = operands_in_alternate_memory_map_.find(instr); - if (operand_it != operands_in_alternate_memory_map_.end()) { - continue; - } - - // If the instruction is a successful async conversion, we don't run the - // post-allocation transformation. - if (successful_async_conversion_set_.contains(instr)) { - continue; - } - - // If any of the operands of the instruction has an in-place user, we - // don't run the post-allocation transformation. - bool in_place_user = false; - for (HloInstruction* operand : instr->operands()) { - // We don't care about users of constants. - if (operand->opcode() == HloOpcode::kConstant) { - continue; - } - if (has_in_place_user(operand)) { - in_place_user = true; - break; - } - } - if (in_place_user) { - continue; - } - - VLOG(3) << "Running post allocation transformation on: \n" - << instr->ToString(); - ABSL_ASSIGN_OR_RETURN(PostAllocationTransformationUpdate changes, - options_.post_allocation_transformation_fn(instr)); - if (!changes.to_be_removed.empty()) { - VLOG(3) << "Post allocation transformation info: \n" - << changes.ToString(); - FixAllocationSequenceAfterPostAllocationTransformation(allocations_, - changes); - } - } - } - } - - if (options_.IsOpSpanExposureEnabled()) { - CHECK_OK(WindowPrefetch()); - } - - if (options_.expanded_scoped_alternate_memory_mode == - ExpandedScopedAlternateMemoryMode::ENABLED) { - ExtendScopedAlternateMemoryAllocations(); - } - - HeapSimulator::Result result; - result.heap_size = result_.heap_size; - result.heap_results.emplace_back(std::move(result_)); - return result; -} - -namespace { - -// Convert a tuple HloUse to its equivalent HloPosition. -HloPosition TupleUseToPosition(const HloUse& use) { - CHECK_EQ(use.instruction->opcode(), HloOpcode::kTuple); - ShapeIndex index = use.operand_index; - index.push_front(use.operand_number); - return {use.instruction, index}; -} - -// Returns the memory space of the defining position of an Allocation object. -MemorySpace GetDefiningPositionMemorySpace(const Allocation& allocation) { - if (!allocation.is_copy_like_allocation()) { - return allocation.memory_space(); - } - if (allocation.memory_space() == MemorySpace::kDefault) { - return MemorySpace::kAlternate; - } - return MemorySpace::kDefault; -} - -} // namespace - -std::vector> -MsaAlgorithm::GetLinkedAllocationsInAlternateMemory( - absl::Span allocation_values) const { - std::vector> linked_allocations; - // A map from position to index into linked_allocations. - absl::flat_hash_map link_id_map; - // Iterate over the allocation values. Find Allocation objects across the - // allocation values that are part of the same linked allocation group. We - // define a linked allocation group as Allocation objects that have aliased - // positions or uses. An example would be an Allocation object that has a - // dynamic-update-slice use and another Allocation object that has the same - // dynamic-update-slice as its defining position. - for (const AllocationValue& allocation_value : allocation_values) { - absl::flat_hash_map> aliases; - for (const AllocationValue::Use& allocation_value_use : - allocation_value.uses()) { - if (!allocation_value_use.aliases.empty()) { - aliases[allocation_value_use.hlo_use] = allocation_value_use.aliases; - } - } - for (const auto& allocation : *allocation_value.allocation_sequence()) { - MemorySpace position_memory_space = - GetDefiningPositionMemorySpace(*allocation); - if (allocation->memory_space() == MemorySpace::kDefault && - position_memory_space == MemorySpace::kDefault) { - // This is just a regular allocation in the default memory, skip. - continue; - } - int link_id = -1; - // For every position and use in the alternate memory space, check if - // there is already a linked allocation group, and if so, use that link - // index. - if (position_memory_space == MemorySpace::kAlternate) { - auto link_id_map_it = link_id_map.find(allocation->defining_position()); - if (link_id_map_it != link_id_map.end()) { - link_id = link_id_map_it->second; - } - } - if (allocation->memory_space() == MemorySpace::kAlternate) { - for (const HloUse& use : allocation->uses()) { - if (use.instruction->opcode() == HloOpcode::kTuple) { - auto link_id_map_it = link_id_map.find(TupleUseToPosition(use)); - if (link_id_map_it != link_id_map.end()) { - if (link_id != -1 && link_id != link_id_map_it->second) { - // We found multiple link indices for the given allocation. We - // merge the two linked allocation groups in that case. - int old_link_id = link_id_map_it->second; - if (old_link_id < link_id) { - std::swap(link_id, old_link_id); - } - absl::c_copy(linked_allocations[old_link_id], - std::back_inserter(linked_allocations[link_id])); - linked_allocations[old_link_id].clear(); - for (auto it = link_id_map.begin(); it != link_id_map.end(); - ++it) { - if (it->second == old_link_id) { - it->second = link_id; - } - } - } - link_id = link_id_map_it->second; - } - } - } - } - if (link_id == -1) { - // Create a new linked allocation group if we couldn't find one. - link_id = linked_allocations.size(); - linked_allocations.push_back({allocation.get()}); - } else { - linked_allocations[link_id].push_back(allocation.get()); - } - // Propagate the link index to all of the aliases of uses in the alternate - // memory. - if (allocation->memory_space() == MemorySpace::kAlternate) { - for (const HloUse& use : allocation->uses()) { - auto alias_it = aliases.find(use); - if (alias_it != aliases.end()) { - for (const HloPosition& aliased_position : alias_it->second) { - link_id_map[aliased_position] = link_id; - } - } - } - } - } - } - - linked_allocations.erase( - std::remove_if( - linked_allocations.begin(), linked_allocations.end(), - [](const auto& allocations) { return allocations.empty(); }), - linked_allocations.end()); - - if (VLOG_IS_ON(3)) { - for (int i = 0; i < linked_allocations.size(); ++i) { - VLOG(3) << " Link id = " << i; - for (const Allocation* allocation : linked_allocations[i]) { - VLOG(3) << " " << allocation->ToString(); - } - } - } - return linked_allocations; -} - -std::vector MsaAlgorithm::GetInefficientAllocationSites( - absl::Span allocation_values) const { - // The logic below is used mostly for testing, allowing a test case to inject - // some custom logic for this method. - if (options_.get_inefficient_allocation_sites_fn) { - std::vector positions; - positions.reserve(allocation_values.size()); - for (const AllocationValue& value : allocation_values) { - positions.push_back(value.position()); - } - return options_.get_inefficient_allocation_sites_fn( - absl::MakeSpan(positions)); - } - - if (!options_.cost_analysis || - options_.inefficient_use_to_copy_ratio == 0.0) { - return {}; - } - - int64_t size = 0; - if (!allocation_values.empty()) { - size = allocation_values.at(0).size(); - } - - if (VLOG_IS_ON(3)) { - VLOG(3) << "Inefficient allocation calculations:"; - for (const AllocationValue& allocation_value : allocation_values) { - for (const auto& allocation : *allocation_value.allocation_sequence()) { - VLOG(3) << " Allocation: " << allocation->ToString(); - if (!allocation->is_copy_like_allocation()) { - const HloPosition& defining_position = - allocation->defining_position(); - int64_t accessed = options_.cost_analysis->OutputBytesAccessed( - *defining_position.instruction, defining_position.index); - VLOG(3) << " pos: " << defining_position.ToString() - << ", accessed: " << accessed << " / " << size; - } - for (const HloUse& use : allocation->uses()) { - int64_t accessed = options_.cost_analysis->OperandBytesAccessed( - *use.instruction, use.operand_number, use.operand_index); - VLOG(3) << " use: " << use.ToString() << ", accessed: " << accessed - << " / " << size; - } - } - } - } - - std::vector> linked_allocations = - GetLinkedAllocationsInAlternateMemory(allocation_values); - std::vector inefficient_sites; - for (const std::vector& allocation_group : - linked_allocations) { - // For all of allocation in the linked allocation group, calculate the total - // use bytes in alternate memory and async copy bytes. If the ratio between - // the two is below inefficient_use_to_copy_ratio, add all of the - // participating allocation sites into inefficient_sites. - VLOG(3) << " AllocationGroup:"; - int64_t copy_bytes = 0; - int64_t use_bytes = 0; - for (const Allocation* allocation : allocation_group) { - VLOG(3) << " Allocation: " << allocation->ToString(); - MemorySpace position_memory_space = - GetDefiningPositionMemorySpace(*allocation); - if (allocation->is_copy_like_allocation()) { - copy_bytes += size; - } - if (position_memory_space == MemorySpace::kAlternate) { - use_bytes += options_.cost_analysis->OutputBytesAccessed( - *allocation->defining_position().instruction, - allocation->defining_position().index); - } - if (allocation->memory_space() == MemorySpace::kAlternate) { - for (const HloUse& use : allocation->uses()) { - use_bytes += options_.cost_analysis->OperandBytesAccessed( - *use.instruction, use.operand_number, use.operand_index); - } - } - } - use_bytes = std::max(use_bytes, int64_t{0}); - copy_bytes = std::max(copy_bytes, int64_t{0}); - VLOG(3) << " use bytes: " << use_bytes - << ", copy bytes: " << copy_bytes; - if (copy_bytes > 0 && - options_.inefficient_use_to_copy_ratio * copy_bytes > use_bytes) { - for (const Allocation* allocation : allocation_group) { - MemorySpace position_memory_space = - GetDefiningPositionMemorySpace(*allocation); - if (position_memory_space == MemorySpace::kAlternate) { - if (!allocation->is_copy_like_allocation()) { - inefficient_sites.push_back(allocation->defining_position()); - } - } - if (allocation->memory_space() == MemorySpace::kAlternate) { - for (const HloUse& use : allocation->uses()) { - inefficient_sites.push_back(use); - } - } - } - } - } - return inefficient_sites; -} - -void MsaAlgorithm::AddRequiredAssignmentsForConditionalOutputsIfNecessary( - absl::Span colocated_intervals) { - for (const MsaBufferInterval* colocated_interval : colocated_intervals) { - const HloValue* value = colocated_interval->buffer; - for (const auto& position : value->positions()) { - if (position.instruction->opcode() == HloOpcode::kConditional && - RequireConditionalOutputsInDefaultMemory(position, value)) { - VLOG(3) << "Adding required assignment for condition output: " - << value->ToShortString(); - { - AddRequiredAssignment( - position.instruction, position.index, MemorySpace::kDefault, - RequiredMemoryAssignment::Source::kConditionalPhiOutput); - } - for (const HloComputation* called_computation : - position.instruction->called_computations()) { - AddRequiredAssignment( - called_computation->root_instruction(), position.index, - MemorySpace::kDefault, - RequiredMemoryAssignment::Source::kConditionalComputationOutput); - } - } - } - } -} - -void MsaAlgorithm::CreateAllocationValuesFromColocatedIntervals( - absl::Span colocated_intervals, - std::vector& allocation_values) { - std::vector new_allocation_values; - // Create AllocationValues for all the colocated intervals. - for (const MsaBufferInterval* colocated_interval : colocated_intervals) { - CreateAllocationValues(*colocated_interval, new_allocation_values); - } - // Go through the AllocationValues and delete the ones that have the identical - // defining instruction and use instructions. This is useful for async - // operations that can read and write to the same buffer, e.g., in-place - // asynchronous collective permute. The AllocationValues that corresponds to - // collective-permute-start{0} (the input) and collective-permute-start{1} - // (the output) refer to the same buffer by definition (since they are created - // from colocated intervals). If we don't delete one of these buffers, then - // when we try to allocate the AllocationValue, we would think they overlap. - auto create_instruction_vector = [](const AllocationValue& allocation_value) { - std::vector instruction_vector; - instruction_vector.push_back(allocation_value.position().instruction); - for (const AllocationValue::Use& use : allocation_value.uses()) { - instruction_vector.push_back(use.hlo_use.instruction); - } - return instruction_vector; - }; - for (int i = 0; i < new_allocation_values.size() - 1; ++i) { - for (int j = i + 1; j < new_allocation_values.size(); ++j) { - const AllocationValue& allocation_value_1 = new_allocation_values[i]; - const AllocationValue& allocation_value_2 = new_allocation_values[j]; - if (allocation_value_1.position().index == - allocation_value_2.position().index && - create_instruction_vector(allocation_value_1) == - create_instruction_vector(allocation_value_2)) { - VLOG(3) << "Allocation values " << allocation_value_1.ToShortString() - << " and " << allocation_value_2.ToShortString() - << " are equivalent, deleting the second one."; - new_allocation_values.erase(new_allocation_values.begin() + j); - --j; - } - } - } - - FindAliases(&new_allocation_values, has_async_pipelined_while_loops_); - absl::c_move(new_allocation_values, std::back_inserter(allocation_values)); -} - -void MsaAlgorithm::MaybeSplitAllocationValues( - absl::Span allocation_values) { - if (options_.determine_split_dimension_fn == nullptr || - options_.shape_size_fn == nullptr || - options_.init_split_tree_fn == nullptr) { - return; - } - - std::vector> results; - - for (AllocationValue& allocation_value : allocation_values) { - std::optional 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]; - auto [it, inserted] = - instruction_to_split_dims_.try_emplace(defining_instruction); - if (inserted) { - it->second = options_.init_split_tree_fn(defining_instruction, nullptr); - } - auto& split_tree = it->second; - const ShapeIndex& shape_index = - allocation_value.value()->defining_position().index; - if (!result.has_value()) { - split_tree[shape_index] = 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. - split_tree[shape_index] = 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 { - if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, allocation_value.position(), - allocation_value.size())) { - return true; - } - return allocation_value.value()->shape().has_layout() && - allocation_value.value()->shape().layout().memory_space() == - options_.alternate_memory_space; -} - -void MsaAlgorithm::AssignDefaultMemIfNotAllowedInAlternateMem( - AllocationValue& allocation_value, int64_t definition_time) { - if (!options_.is_position_allowed_in_alternate_mem_fn( - allocation_value.position()) || - MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - options_.msa_tensor_overrides, allocation_value.position(), - allocation_value.size())) { - std::optional existing_req = - RequiredMemoryAssignmentAt(allocation_value.value(), definition_time); - // If a value is pre-colored or already possesses an explicit alternate - // memory requirement, preserve that requirement and log a warning instead - // of forcing an incompatible assignment in default memory. - if (RequiresNoCopyAlternateMemAllocation(allocation_value) || - IsPositionColoredInAlternateMemory(allocation_value.position()) || - IsPositionColoredInAlternateMemoryAtTime(allocation_value.position(), - definition_time) || - (existing_req.has_value() && - existing_req->memory_space == MemorySpace::kAlternate)) { - LOG(WARNING) << "The value " << allocation_value.value()->ToShortString() - << " is pre-colored for alternate memory but the position " - << allocation_value.position().ToString() - << " is not allowed in the alternate memory. Respecting the " - "color " - "but this may break things later in compilation."; - } else { - AddRequiredAssignment(allocation_value.value(), - allocation_value.position().instruction, - static_cast(MemorySpace::kDefault), - static_cast(definition_time), - RequiredMemoryAssignment::Source:: - kPositionNotAllowedInAlternateMemory); - } - } -} - -std::vector -MsaAlgorithm::GenerateAllocationSegmentContexts( - absl::Span& allocation_values, - absl::flat_hash_map>& - value_indices_by_sync_inst, - int allocation_value_idx) const { - AllocationValue& allocation_value = - allocation_values.at(allocation_value_idx); - std::vector uses_work_list; - for (int primary_use_idx = 0; - primary_use_idx < allocation_value.uses().size(); ++primary_use_idx) { - AllocationValue::Use& primary_use = - allocation_value.uses().at(primary_use_idx); - if (!IsInstructionPendingReplacements(primary_use.hlo_use.instruction)) { - uses_work_list.push_back({&allocation_value.uses(), primary_use_idx, - allocation_value_idx, false}); - } else { - uses_work_list.push_back({&allocation_value.uses(), primary_use_idx, - allocation_value_idx, true}); - for (auto sync_destination_idx : - value_indices_by_sync_inst.at(primary_use.hlo_use.instruction)) { - AllocationValue& sync_destination = - allocation_values.at(sync_destination_idx); - if (sync_destination.position().instruction == - primary_use.hlo_use.instruction) { - VLOG(3) << "Adding secondary uses related to allocation value " - << sync_destination.ToShortString() - << " to uses worklist, because the allocation value is " - "defined at the copy use instruction output."; - for (int secondary_use_id = 0; - secondary_use_id < sync_destination.uses().size(); - ++secondary_use_id) { - // This is an important line - sync_destination.uses().at(secondary_use_id).sync_mem_op_operand = - primary_use.hlo_use.instruction; - int allocation_value_to_update_idx = sync_destination_idx; - uses_work_list.push_back({&sync_destination.uses(), - secondary_use_id, - allocation_value_to_update_idx, false}); - } - } else { - VLOG(3) << "Skipping secondary uses related to allocation value " - << sync_destination.ToShortString() - << ", because the allocation value is not defined at the " - "copy use instruction " - "output."; - } - } - } - } - // Sort uses according to their use time - std::sort(uses_work_list.begin(), uses_work_list.end(), - [](const auto& a, const auto& b) { - return a.uses->at(a.use_idx).time < b.uses->at(b.use_idx).time; - }); - VLOG(3) << "Uses work list:"; - for (int i = 0; i < uses_work_list.size(); i++) { - auto [uses, use_idx, allocation_value_to_update_idx, - only_extend_existing_allocation] = uses_work_list.at(i); - VLOG(3) << " " << i + 1 << "/" << uses_work_list.size() << ") " - << uses->at(use_idx).hlo_use.ToString(); - } - if (uses_work_list.empty()) { - VLOG(3) << " --Empty--"; - } - return uses_work_list; -} - -bool MsaAlgorithm::RequireConditionalOutputsInDefaultMemory( - HloPosition conditional_phi_position, const HloValue* hlo_value) { - CHECK(conditional_phi_position.instruction->opcode() == - HloOpcode::kConditional); - - // Check if the phi is required to be in the default memory. - std::optional required_assignment_at_definition = - RequiredMemoryAssignmentAt(hlo_value, - hlo_live_range_.instruction_schedule().at( - conditional_phi_position.instruction)); - if (required_assignment_at_definition.has_value() && - required_assignment_at_definition->memory_space == - MemorySpace::kDefault) { - return true; - } - - // Check if the branched computation roots are not tuples or if they are - // required to be in the default memory. - for (const HloComputation* branched_computation : - conditional_phi_position.instruction->called_computations()) { - // If the branched computation root is not a tuple, then the output - // allocation value inside the branched computation has no uses and cannot - // be placed in the alternate memory. - if (branched_computation->root_instruction()->opcode() != - HloOpcode::kTuple) { - return true; - } - - // Check if the branched computation root value is required to be in the - // default memory at the root instruction. - const HloValue& computation_root_value = - alias_analysis_.dataflow_analysis().GetUniqueValueAt( - branched_computation->root_instruction(), - conditional_phi_position.index); - std::optional required_assignment_in_branch = - RequiredMemoryAssignmentAt( - &computation_root_value, - hlo_live_range_.instruction_schedule().at( - branched_computation->root_instruction())); - if (required_assignment_in_branch.has_value() && - required_assignment_in_branch->memory_space == MemorySpace::kDefault) { - return true; - } - } - return false; -} - -void MsaAlgorithm::AddOperandToAlternateMemoryMap( - const HloInstruction* instruction, int operand_number, - const ShapeIndex& index) { - operands_in_alternate_memory_map_[instruction].insert( - std::make_pair(operand_number, index)); -} - -bool MsaAlgorithm::GetUpdatedRequireNoCopyAlternateMemForAsyncPipelinedWhile( - bool require_no_copy_alt_mem, - const AllocationValue& allocation_value_to_update, - const AliasedOffset* preferred_offset, - const absl::flat_hash_set& default_req_buffer_ids) const { - // Check 1: leave normal MSA as is. - if (!has_async_pipelined_while_loops_) { - return require_no_copy_alt_mem; - } - if (!IsBufferAliasedToAsyncPipelinedWhileLoop( - allocation_value_to_update.value())) { - return require_no_copy_alt_mem; - } - // Check 2: we're only looking to override false -> true. - if (require_no_copy_alt_mem) { - return true; - } - - // Check 3: we're only looking to override async dynamic-update-slice update - // slice and async dynamic-slice output. - const HloPosition& def_pos = allocation_value_to_update.position(); - if (!IsAsyncPipelinedWhileAlternateMemoryPosition(def_pos, alias_analysis_)) { - return false; - } - - if (preferred_offset == nullptr) { - return false; - } - - if (!options_.is_position_allowed_in_alternate_mem_fn(def_pos) || - options_.size_fn(*allocation_value_to_update.value()) > - options_.max_size_in_bytes) { - return false; - } - - int64_t buffer_id = - alias_analysis_ - .GetUniqueBufferAt( - allocation_value_to_update.value()->defining_position()) - .id(); - return !default_req_buffer_ids.contains(buffer_id); -} - -absl::StatusOr MsaAlgorithm::AllocateAllocationValues( - absl::Span allocation_values) { - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - absl::flat_hash_map> - value_indices_by_sync_inst; - for (size_t idx = 0; idx < allocation_values.size(); ++idx) { - const HloInstruction* inst = - allocation_values.at(idx).position().instruction; - if (IsInstructionPendingReplacements(inst)) { - value_indices_by_sync_inst[inst].push_back(idx); - } - } - - // Extract all use times. Use the view-extended time so this list agrees with - // the chunk reservation end (request.end_time) computed in - // CreateAllocationRequest; otherwise a view-extended end_time would not be - // found in this list. - std::vector all_use_times; - for (const AllocationValue& allocation_value : allocation_values) { - for (const auto& use : allocation_value.uses()) { - all_use_times.push_back(GetExtendedUseTimeIfUseIsView(use.hlo_use)); - } - } - absl::c_sort(all_use_times); - for (int i = 0; i < all_use_times.size(); ++i) { - 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. - absl::flat_hash_map - preferred_offset_for_computation; - absl::flat_hash_map - preferred_offset_for_allocation_value; - absl::flat_hash_map - definition_time_for_allocation_value; - - // These are allocation values inside a conditional that have already been - // processed. They may get an allocation in the alternate memory that mirrors - // a real allocation outside the conditional, that is live throughout the - // conditional. - absl::flat_hash_set - already_processed_allocation_values_inside_a_conditional; - - absl::flat_hash_set default_req_buffer_ids; - if (has_async_pipelined_while_loops_) { - // NOLINTNEXTLINE - for (const auto& [val, req_list] : required_assignments_) { - for (const auto& req : req_list) { - if (req.memory_space == MemorySpace::kDefault) { - default_req_buffer_ids.insert( - alias_analysis_.GetUniqueBufferAt(val->defining_position()).id()); - break; - } - } - } - } - - AllocationResult result = AllocationResult::kSuccess; - for (int alloc_value_idx = 0; alloc_value_idx < allocation_values.size(); - ++alloc_value_idx) { - auto& allocation_value = allocation_values.at(alloc_value_idx); - VLOG(3) << alloc_value_idx + 1 << "/" << allocation_values.size() - << ") Allocating allocation value: " - << allocation_value.ToShortString(); - - if (IsInstructionPendingReplacements( - allocation_value.position().instruction)) { - VLOG(3) << "Skip allocating allocation value " - << allocation_value.ToShortString(); - continue; - } - - VLOG(4) << "require_no_copy_alternate_mem_allocation = " - << RequiresNoCopyAlternateMemAllocation(allocation_value); - if (RequiresNoCopyAlternateMemAllocation(allocation_value) && - allocation_value.size() > available_heap_size()) { - if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, allocation_value.position(), - allocation_value.size())) { - return absl::ResourceExhaustedError(absl::StrCat( - "Cannot allocate pinned tensor in alternate memory: tensor size (", - allocation_value.size(), " bytes) exceeds available heap size (", - available_heap_size(), " bytes) for defining instruction '", - allocation_value.position().instruction->name(), "'")); - } - VLOG(3) << "Skip " << allocation_value.value()->ToShortString() - << " because the buffer is larger than the heap size."; - continue; - } - - const AllocationValue::Use* previous_use = nullptr; - auto uses_work_list = GenerateAllocationSegmentContexts( - allocation_values, value_indices_by_sync_inst, alloc_value_idx); - - // Iterate over the uses. - for (auto& entry : uses_work_list) { - const AllocationValue::Use& use = entry.uses->at(entry.use_idx); - AllocationValue& allocation_value_to_update = - allocation_values.at(entry.allocation_value_to_update_idx); - std::string extension_only_hint_str = - entry.only_extend_existing_allocation ? " (extension only): " : ": "; - VLOG(3) << "Working on use" << extension_only_hint_str - << use.hlo_use.ToString() - << ", allocation value: " << allocation_value.ToShortString() - << ", updates allocation value: " - << allocation_value_to_update.ToShortString(); - - if (!definition_time_for_allocation_value.contains( - &allocation_value_to_update)) { - definition_time_for_allocation_value[&allocation_value_to_update] = - hlo_live_range_.instruction_schedule().at( - allocation_value_to_update.position().instruction); - AssignDefaultMemIfNotAllowedInAlternateMem( - allocation_value_to_update, definition_time_for_allocation_value.at( - &allocation_value_to_update)); - } - - if (!preferred_offset_for_allocation_value.contains( - &allocation_value_to_update)) { - // Resolve the preferred alternate memory offset for this allocation - // value by inspecting the canonical HloBuffer alias set. This - // guarantees that all aliased positions across while loop boundaries - // (including loop input operands, loop parameters, and loop body root - // tuple elements) adopt identical starting offsets in alternate memory. - AliasedOffset* preferred_offset = nullptr; - if (has_async_pipelined_while_loops_) { - const HloBuffer& hlo_buffer = alias_analysis_.GetUniqueBufferAt( - allocation_value_to_update.position().instruction, - allocation_value_to_update.position().index); - auto buf_it = pipelined_while_buffer_id_to_aliased_offset_.find( - hlo_buffer.id()); - if (buf_it != pipelined_while_buffer_id_to_aliased_offset_.end()) { - preferred_offset = buf_it->second; - } else { - for (const HloPosition& position : hlo_buffer.ComputePositions()) { - if (!IsAsyncPipelinedWhilePosition(position)) { - continue; - } - HloComputation* comp = position.instruction->parent(); - auto comp_it = - pipelined_while_preferred_offset_for_computation_.find(comp); - if (comp_it != - pipelined_while_preferred_offset_for_computation_.end()) { - auto index_it = comp_it->second.find(position.index); - if (index_it != comp_it->second.end()) { - preferred_offset = index_it->second; - break; - } - } - } - } - } - if (preferred_offset == nullptr) { - auto comp_it = preferred_offset_for_computation.find( - allocation_value_to_update.computation()); - if (comp_it != preferred_offset_for_computation.end()) { - preferred_offset = comp_it->second; - } else if (has_async_pipelined_while_loops_ && - IsAsyncPipelinedWhilePosition( - allocation_value_to_update.position())) { - auto comp_it = - pipelined_while_preferred_offset_for_computation_.find( - allocation_value_to_update.computation()); - if (comp_it != - pipelined_while_preferred_offset_for_computation_.end()) { - auto index_it = comp_it->second.find( - allocation_value_to_update.position().index); - if (index_it != comp_it->second.end()) { - preferred_offset = index_it->second; - } - } - } - } - preferred_offset_for_allocation_value[&allocation_value_to_update] = - preferred_offset; - } - preferred_offset_for_allocation_value[&allocation_value_to_update] = - CheckOrUpdatePreferredOffsetForUse( - use, preferred_offset_for_allocation_value.at( - &allocation_value_to_update)); - bool require_no_copy_alt_mem = - RequiresNoCopyAlternateMemAllocation(allocation_value_to_update); - require_no_copy_alt_mem = - GetUpdatedRequireNoCopyAlternateMemForAsyncPipelinedWhile( - require_no_copy_alt_mem, allocation_value_to_update, - preferred_offset_for_allocation_value.at( - &allocation_value_to_update), - default_req_buffer_ids); - AllocationRequest request = CreateAllocationRequest( - allocation_value, allocation_value_to_update, use, previous_use, - preferred_offset_for_allocation_value.at(&allocation_value_to_update), - definition_time_for_allocation_value.at(&allocation_value_to_update), - require_no_copy_alt_mem, all_use_times, - entry.only_extend_existing_allocation, - allocation_values.subspan(0, alloc_value_idx), allocation_values, - /*shape_override=*/std::nullopt); - if (options_.allocation_request_modifier_testing_fn) { - options_.allocation_request_modifier_testing_fn(request); - } - bool is_processed_allocation_value_live_throughout_a_conditional = - already_processed_allocation_values_inside_a_conditional.contains( - &allocation_value); - // Skip trivial uses that are not the root instruction. - if (!IsUseTrivialNonRoot(use.hlo_use) && - !is_processed_allocation_value_live_throughout_a_conditional) { - UpdateRequestWithAlternateMemoryColoringRequirements(request); - UpdateRequestWithDefaultMemoryColoringRequirements(request); - AllocationResult allocate_segment_result = AllocateSegment(request); - VLOG(2) << "AllocateSegment result: " - << ResultToString(allocate_segment_result); - result_mark(allocate_segment_result, result); - AllocationSequence* allocation_sequence = - allocation_value_to_update.mutable_allocation_sequence(); - if (options_.allocation_result_modifier_testing_fn) { - options_.allocation_result_modifier_testing_fn( - request, result, - options_.prefetch_interval_picker->retry_number()); - } - - if (request.fail_on_unsatisfied_override || request.strict_timing || - (request.require_no_copy_alternate_mem_allocation && - MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, - allocation_value_to_update.position(), - allocation_value_to_update.size()))) { - if (allocate_segment_result != AllocationResult::kSuccess) { - std::string reason = ResultToString(allocate_segment_result); - if (result_is(allocate_segment_result, - AllocationResult::kFailOutOfMemory)) { - reason = "Out of alternate memory / spatial placement conflict"; - } else if (result_is( - allocate_segment_result, - AllocationResult::kFailViolatesAsyncCopyResource)) { - reason = "Copy resource limit reached / bandwidth saturation"; - } else if (result_is(allocate_segment_result, - AllocationResult::kFailOutOfAsyncCopies)) { - reason = "Ran out of outstanding asynchronous copies"; - } else if (result_is(allocate_segment_result, - AllocationResult::kFailLiveRangeTooShort)) { - reason = "Live range too short / dependency ordering"; - } - return absl::FailedPreconditionError(absl::StrCat( - "MSA strict override failed for instruction '", - request.use->hlo_use.instruction->name(), "' operand ", - request.use->hlo_use.operand_number, ", shape ", - request.use->hlo_use.instruction - ->operand(request.use->hlo_use.operand_number) - ->shape() - .ToString(), - ", defining instruction '", - allocation_value_to_update.position().instruction->name(), - "': ", reason, " (requested schedule time: ", - request.preferred_prefetch_time.has_value() - ? absl::StrCat(*request.preferred_prefetch_time) - : "n/a", - ")")); - } - if (allocate_segment_result == AllocationResult::kSuccess && - request.strict_timing && - request.preferred_prefetch_time.has_value()) { - const Allocation* allocation = - allocation_sequence->empty() - ? nullptr - : allocation_sequence->back().get(); - int64_t actual_prefetch_time = 0; - if (allocation && allocation->is_copy_allocation()) { - actual_prefetch_time = - static_cast(allocation) - ->copy_start_schedule_after(); - } else if (allocation && allocation->is_sliced_copy_allocation()) { - actual_prefetch_time = - static_cast(allocation) - ->slice_details_sorted_by_start_time() - .front() - .copy_start_after_time; - } - if (actual_prefetch_time != *request.preferred_prefetch_time) { - return absl::FailedPreconditionError(absl::StrCat( - "MSA strict timing override unsatisfied for instruction '", - request.use->hlo_use.instruction->name(), "' operand ", - request.use->hlo_use.operand_number, ", shape ", - request.use->hlo_use.instruction - ->operand(request.use->hlo_use.operand_number) - ->shape() - .ToString(), - ": scheduled prefetch time (", actual_prefetch_time, - ") does not match requested prefetch time (", - *request.preferred_prefetch_time, ")")); - } - } - } - if (allocate_segment_result == AllocationResult::kSuccess && - ShouldBeMirrored(allocation_value_to_update, use, previous_use)) { - CreateMirroredAllocations( - allocation_value_to_update, use, previous_use, allocation_values, - already_processed_allocation_values_inside_a_conditional); - } - if (request.require_copy_allocation) { - auto it = std::find_if( - allocation_sequence->begin(), allocation_sequence->end(), - [&](const std::unique_ptr< - xla::memory_space_assignment::Allocation>& allocation_ptr) { - if (allocation_ptr->is_copy_allocation()) { - auto copy_allocation = - dynamic_cast(allocation_ptr.get()); - return copy_allocation && - (copy_allocation->copy_done_schedule_before() <= - request.required_copy_allocation_latest_time) && - (copy_allocation->sync_mem_op() == - request.required_copy_allocation_for) && - (!request.required_copy_for_slice || - (request.required_copy_for_slice && - !copy_allocation->cross_program_prefetch_index() - .has_value())); - } - if (allocation_ptr->is_sliced_copy_allocation()) { - auto sliced_copy_allocation = - dynamic_cast( - allocation_ptr.get()); - return sliced_copy_allocation && - (sliced_copy_allocation->earliest_available_time() <= - request.required_copy_allocation_latest_time) && - (sliced_copy_allocation->sync_mem_op() == - request.required_copy_allocation_for) && - !request.required_copy_for_slice; - } - return false; - }); - - if (result_requires_uncommit(result) || - it == allocation_sequence->end()) { - VLOG(3) << "No async copy allocation found by the end of " - "segment allocation. " - "Sync copy replacement has failed. Fall back to the " - "normal mode."; - failed_async_conversions_[request.required_copy_allocation_for] = - AsyncConversionResult::kFailedSatisfyingConstraints; - result_mark(AllocationResult::kFailSyncDataMoveReplacement, result); - result_mark(AllocationResult::kFailRequiresUncommit, result); - } else { - bool has_correct_use = false; - for (auto& alloc_use : (*it)->uses()) { - if (alloc_use == request.use->hlo_use) { - has_correct_use = true; - break; - } - } - if (!has_correct_use) { - VLOG(3) << "No async copy allocation found by the end of " - "segment allocation with the correct use. " - "Sync copy replacement has failed. Fall back to the " - "normal mode."; - failed_async_conversions_[request.required_copy_allocation_for] = - AsyncConversionResult::kFailedPrecondition; - result_mark(AllocationResult::kFailSyncDataMoveReplacement, - result); - result_mark(AllocationResult::kFailRequiresUncommit, result); - } else { - not_finalized_async_conversions_.push_back( - request.required_copy_allocation_for); - VLOG(3) << "Replacing " - << request.required_copy_allocation_for->ToShortString() - << " with " << (*it)->ToString(); - } - } - } - if (request.require_no_copy_alternate_mem_allocation && - result != AllocationResult::kSuccess) { - absl::Status failed_precondition = FailedPrecondition( - "The value defined at %s requires allocation in the alternate " - "memory, which could not be satisfied. This typically happens " - "because more pinned buffers are live than the alternate memory " - "capacity.", - allocation_value.position().ToString()); - LOG(ERROR) << failed_precondition; - return failed_precondition; - } - if (result_requires_uncommit(result)) { - // If the allocation finding failed (e.g., due to running out of - // asynchronous copies), then fall back to allocating the buffer - // entirely in the default memory. - return result; - } - - // If there are multiple uses, they can try using the memory - // allocation already at the alternate memory. - definition_time_for_allocation_value[&allocation_value_to_update] = - instruction_schedule.at(use.hlo_use.instruction); - previous_use = &use; - } - if (entry.only_extend_existing_allocation) { - continue; - } - const auto use_time = request.end_time; - if (!TryUpdateAllocationRequirementForUseAliases( - allocation_value_to_update, use, use_time)) { - // The use's allocation contradicts an existing required assignment at - // an aliased position (e.g. among positions aliased by an in-place - // instruction such as a dynamic-update-slice, this use's allocation - // landed in alternate memory while another aliased use already - // required default memory there). Uncommit and retry; the retry - // honors the existing requirement, ultimately falling back to - // default memory. - result_mark(AllocationResult::kFailRequiresUncommit, result); - return result; - } - MaybeCreateMirroredParentAllocationForWhileUse( - allocation_value_to_update, use, use_time, allocation_values, - preferred_offset_for_computation); - } - } - - if (!VerifyAllConversionsAreSuccessful()) { - result_mark(AllocationResult::kFailSyncDataMoveReplacement, result); - result_mark(AllocationResult::kFailRequiresUncommit, result); - } - - return result; -} - -bool MsaAlgorithm::VerifyAllConversionsAreSuccessful() { - for (const HloInstruction* instruction : - sorted_async_conversion_candidates_) { - if (absl::c_find(not_finalized_async_conversions_, instruction) == - not_finalized_async_conversions_.end()) { - if (!failed_async_conversions_.contains(instruction)) { - failed_async_conversions_[instruction] = - AsyncConversionResult::kFailedNotProcessed; - VLOG(3) << "Async conversion failed for " - << instruction->ToShortString() - << " because its operand or user was not processed."; - } - return false; - } - } - return true; -} - -AliasedOffset* MsaAlgorithm::CheckOrUpdatePreferredOffsetForUse( - const AllocationValue::Use& use, AliasedOffset* preferred_offset) const { - // Assign the required assignment offset as a preferred offset. - std::optional required_assignment = - AliasedRequiredAssignmentForUse(use); - if (required_assignment.has_value() && - required_assignment.value().memory_space == MemorySpace::kAlternate) { - if (preferred_offset && preferred_offset->offset != - required_assignment.value().offset->offset) { - VLOG(2) << "Overriding preferred offset " << preferred_offset->offset - << " with required assignment offset " - << required_assignment.value().offset->offset << " for use:\n" - << use.hlo_use.ToString(); - } - preferred_offset = required_assignment->offset; - } - return preferred_offset; -} - -AllocationRequest MsaAlgorithm::CreateAllocationRequest( - AllocationValue& allocation_value, - AllocationValue& allocation_value_to_update, - const AllocationValue::Use& use, const AllocationValue::Use* previous_use, - AliasedOffset* preferred_offset, int64_t definition_time, - bool require_no_copy_alternate_mem_allocation, - const std::vector& all_use_times, - bool only_extend_existing_allocation, - absl::Span processed_allocation_values, - absl::Span all_allocation_values, - std::optional shape_override) { - const HloUse& hlo_use = use.hlo_use; - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - bool require_copy_allocation = false; - int64_t required_copy_allocation_latest_time = 0; - HloInstruction* required_copy_allocation_for = nullptr; - bool required_copy_for_slice = false; - std::optional earliest_prefetch_time = std::nullopt; - if (use.sync_mem_op_operand && - IsInstructionPendingReplacements(use.sync_mem_op_operand)) { - required_copy_allocation_for = use.sync_mem_op_operand; - require_copy_allocation = true; - required_copy_for_slice = - (IsAsyncConversionSliceCandidate(use.sync_mem_op_operand) == - AsyncConversionResult::kSuccess); - // The async copy allocation can be delayed until the earliest time at which - // the value is used in a position or the earliest use time of the updated - // allocation value. We find the minimum of these two times. - int64_t min_time = - GetCorrectedUseTime(allocation_value.position().instruction); - int64_t earliest_position_time = std::numeric_limits::max(); - for (auto& position : allocation_value.value()->positions()) { - auto position_time = GetCorrectedUseTime(position.instruction); - if (position_time > min_time) { - earliest_position_time = - std::min(earliest_position_time, position_time); - } - } - int64_t earliest_use_time = std::numeric_limits::max(); - for (auto& secondary_use : allocation_value_to_update.uses()) { - if (!IsTrivialInstruction(secondary_use.hlo_use.instruction) || - secondary_use.hlo_use.instruction == - use.hlo_use.instruction->parent()->root_instruction()) { - earliest_use_time = std::min( - earliest_use_time, GetCorrectedUseTime(secondary_use.hlo_use)); - } - } - required_copy_allocation_latest_time = - std::min(earliest_use_time, earliest_position_time); - // We need to make sure that the copy allocation is scheduled before the - // controlled successor of the sync mem op. - for (const HloInstruction* control_successor : - required_copy_allocation_for->control_successors()) { - auto successor_it = instruction_schedule.find(control_successor); - if (successor_it != instruction_schedule.end()) { - int64_t successor_time = successor_it->second; - if (successor_time < required_copy_allocation_latest_time) { - VLOG(3) - << "Updating the required replacement async mem op allocation " - "latest time from " - << required_copy_allocation_latest_time << " to " - << successor_time << ", because of control successor " - << control_successor->ToString(); - required_copy_allocation_latest_time = successor_time; - } - } - } - - // Make sure that the earliest prefetch time is after the schedule time of - // the latest operand of the sync mem op. - for (const HloInstruction* operand : - required_copy_allocation_for->operands()) { - earliest_prefetch_time = std::max(earliest_prefetch_time.value_or(-1), - instruction_schedule.at(operand)); - } - } - - int64_t use_time = GetExtendedUseTimeIfUseIsView(hlo_use); - bool allow_no_copy_alternate_mem_allocation = true; - bool allow_prefetch = true; - bool prefer_no_copy_alternate_mem_allocation = false; - // TODO(b/318886791): Rename boundary variables (here and other places) - // like `latest_prefetch_time` and `earliest_prefetch_time` indicate - // whether they are exclusive or inclusive boundaries. - // A view use extends `use_time` through the view's transitive readers so - // the buffer stays reserved while they read through the view, but a - // prefetched copy is materialized right before the view instruction - // itself. Bound the prefetch deadline by the actual use time; otherwise - // the copy is reserved over (extended_start, extended_end) while it - // actually runs at the view's position, and the unreserved gap in between - // lets other buffers land on the same offsets (verifier chunk overlap). - int64_t latest_prefetch_time = - hlo_live_range_.instruction_schedule().at(hlo_use.instruction); - - // Control flow calls include kWhile, kCall, and kConditional opcodes. - bool is_sequential_call = - (GetInstructionCallContext(hlo_use.instruction->opcode()) == - CallContext::kControlFlow); - if (is_sequential_call) { - for (const HloComputation* called_computation : - hlo_use.instruction->called_computations()) { - auto called_computation_span_it = - hlo_live_range_.computation_span_times().find(called_computation); - if (called_computation_span_it != - hlo_live_range_.computation_span_times().end()) { - latest_prefetch_time = std::min( - called_computation_span_it->second.start - 1, latest_prefetch_time); - } - } - use_time = GetCorrectedUseTime(hlo_use); - } - - // Add a required assignment in default memory if the use not allowed in - // alternate memory. - std::optional existing_req = - RequiredMemoryAssignmentAt(allocation_value_to_update.value(), use_time); - auto is_required_in_alt_mem = [&]() { - return require_no_copy_alternate_mem_allocation || - IsPositionColoredInAlternateMemory( - allocation_value_to_update.position()) || - IsPositionColoredInAlternateMemoryAtTime( - allocation_value_to_update.position(), use_time) || - IsUseColoredInAlternateMemory(hlo_use) || - (existing_req.has_value() && - existing_req->memory_space == MemorySpace::kAlternate); - }; - - if (IsWhileLoopUseRequiredInDefaultMemory(hlo_use) || - !IsUseAllowedInAlternateMemory(hlo_use) || - !IsWhileLoopUseBeneficialInAlternateMemory(hlo_use)) { - // Respect pre-existing alternate memory coloring or required assignments - // even if a while loop use would not otherwise be allowed or beneficial - // in alternate memory. - if (is_required_in_alt_mem()) { - LOG(WARNING) << "The value " - << allocation_value_to_update.value()->ToShortString() - << " is pre-colored for alternate memory but the use " - << hlo_use.ToString() - << " is not allowed in the alternate memory. Respecting the " - "color but this may break things later in compilation."; - } else { - AddRequiredAssignment( - allocation_value_to_update.value(), hlo_use.instruction, - MemorySpace::kDefault, use_time, - RequiredMemoryAssignment::Source::kUseNotAllowedInAlternateMemory); - } - } else if (!IsConditionalUseBeneficialInAlternateMemory( - allocation_value_to_update, hlo_use)) { - // If there is an existing required assignment for the value in - // alternate memory due to an aliased use, then it is okay to allow this - // use alternate memory, otherwise force default memory allocation. - if (existing_req.has_value() && - existing_req.value().memory_space == MemorySpace::kAlternate && - existing_req.value().required_assignment_source == - RequiredMemoryAssignment::Source::kAliasedUse) { - VLOG(1) << "Allowing use " << hlo_use.ToString() - << " in alternate memory because of existing required" - " assignment due to aliased use: " - << existing_req.value().ToString() << " for value: " - << allocation_value_to_update.value()->ToShortString(); - } else if (is_required_in_alt_mem()) { - VLOG(1) << "Allowing conditional use " << hlo_use.ToString() - << " in alternate memory despite not being beneficial, due to " - "coloring."; - } else { - RequiredMemoryAssignment::Source source = - RequiredMemoryAssignment::Source::kUseNotAllowedInAlternateMemory; - AddRequiredAssignment(allocation_value_to_update.value(), - hlo_use.instruction, MemorySpace::kDefault, - use_time, source); - } - } else if (IsEvictionRequiredForPreviousUseAtConditional( - allocation_value, use, previous_use, all_allocation_values)) { - allow_no_copy_alternate_mem_allocation = false; - auto prev_it = instruction_schedule.find(previous_use->hlo_use.instruction); - if (prev_it != instruction_schedule.end()) { - earliest_prefetch_time = prev_it->second; - } - VLOG(3) << "Previous use (" << previous_use->hlo_use.ToString() - << ") of use (" << hlo_use.ToString() - << ") is a conditional, the operand requires eviction. " - << "Earliest prefetch time = " - << (earliest_prefetch_time.has_value() ? *earliest_prefetch_time - : -1); - } - - AllocationRequest request; - // Bitcasts don't define buffers and don't directly consume buffers. Skip - // allocating buffers for bitcast uses (unless they are the root - // instruction). The uses that feed from bitcasts will be handled - // specially. - if (hlo_use.instruction->opcode() != HloOpcode::kBitcast || - hlo_use.instruction == - hlo_use.instruction->parent()->root_instruction()) { - std::optional preferred_prefetch_time = std::nullopt; - auto loop_optimized_allocation_it = - loop_optimized_allocations_map_.find(use.hlo_use); - if (loop_optimized_allocation_it != loop_optimized_allocations_map_.end()) { - const LoopOptimizedAllocationInfo& loop_optimized_allocation_info = - loop_optimized_allocation_it->second; - const Allocation* allocation = - loop_optimized_allocation_info.loop_optimized_allocation; - VLOG(3) << "Found optimized allocation for " << use.hlo_use.ToString() - << " (loop idx: " << loop_optimized_allocation_info.use_index - << "): " << allocation->ToString(); - if (require_no_copy_alternate_mem_allocation) { - if (allocation->is_copy_allocation() || - allocation->memory_space() == MemorySpace::kDefault) { - VLOG(1) << "Optimized allocation could not be applied " - "because the tensor is pre-colored, allocation: " - << allocation->ToString(); - } - } else if (allocation->is_copy_allocation()) { - allow_no_copy_alternate_mem_allocation = true; - const CopyAllocation* copy_allocation = - static_cast(allocation); - int64_t effective_copy_start_time = - copy_allocation->copy_start_schedule_after(); - if (copy_allocation->copy_start_schedule_after() == - loop_optimized_allocation_info.loop_size - 1 && - copy_allocation->copy_done_schedule_before() == 0) { - effective_copy_start_time = -loop_optimized_allocation_info.loop_size; - } else if (copy_allocation->copy_start_schedule_after() + 1 >= - copy_allocation->copy_done_schedule_before()) { - effective_copy_start_time -= loop_optimized_allocation_info.loop_size; - } - preferred_prefetch_time = use_time - - loop_optimized_allocation_info.use_index + - effective_copy_start_time; - VLOG(3) << "Prefer prefetch at " << *preferred_prefetch_time - << " (effective: " << effective_copy_start_time << ")"; - } else if (allocation->memory_space() == MemorySpace::kDefault) { - allow_prefetch = false; - allow_no_copy_alternate_mem_allocation = false; - VLOG(3) << "Disallowing alternate memory allocation."; - } else { - CHECK(allocation->memory_space() == MemorySpace::kAlternate); - prefer_no_copy_alternate_mem_allocation = true; - VLOG(3) << "Prefer no-copy alternate memory allocation."; - } - } - - // TODO(mehrdadk): Remove this code once we have a better way to find - // repeated instructions. - if (/* DISABLES CODE */ (false)) { - const std::vector* repeated_insts = - GetRepeatedInstructionList(hlo_use.instruction); - if (repeated_insts) { - for (int i = 0; i < repeated_insts->size(); ++i) { - const HloInstruction* repeated = repeated_insts->at(i); - VLOG(4) << "Repeated instruction for use: " << repeated->name() << " " - << hlo_live_range_.instruction_schedule().at(repeated); - if (repeated == hlo_use.instruction && i > 0) { - const HloInstruction* prev_repeated = repeated_insts->at(i - 1); - if (prev_repeated->parent() == hlo_use.instruction->parent()) { - preferred_prefetch_time = - hlo_live_range_.instruction_schedule().at(prev_repeated) + 1; - VLOG(3) << "Found a previous repeated (" << prev_repeated->name() - << ") at " << (*preferred_prefetch_time - 1) - << ". Setting preferred prefetch time = " - << *preferred_prefetch_time; - } - } - } - } - } - - bool strict_timing = false; - bool fail_on_unsatisfied_override = false; - bool is_prefetch_override = false; - bool require_end_colored_in_default_memory = false; - - if (MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - options_.msa_tensor_overrides, hlo_use, allocation_value.size()) || - MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - options_.msa_tensor_overrides, allocation_value.position(), - allocation_value.size())) { - allow_prefetch = false; - allow_no_copy_alternate_mem_allocation = false; - require_end_colored_in_default_memory = true; - AddRequiredAssignment( - allocation_value_to_update.value(), hlo_use.instruction, - MemorySpace::kDefault, use_time, - RequiredMemoryAssignment::Source::kUseNotAllowedInAlternateMemory); - } - if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - options_.msa_tensor_overrides, allocation_value.position(), - allocation_value.size())) { - require_no_copy_alternate_mem_allocation = true; - allow_prefetch = false; - } - - int64_t live_range_start_time = (earliest_prefetch_time.has_value() - ? earliest_prefetch_time.value() - : std::min(definition_time, use_time)); - auto prefetch_override_info = - MemorySpaceAssignmentUtils::GetPrefetchOverrideInfo( - options_.msa_tensor_overrides, - options_.preferred_prefetch_overrides, allocation_value.size(), - hlo_use, instruction_schedule, live_range_start_time, - latest_prefetch_time, allocation_value.position()); - CHECK_OK(prefetch_override_info.status()); - if (prefetch_override_info.value().has_value()) { - VLOG(1) << "Overriding prefetch for " << hlo_use.instruction->name() - << " operand number " << hlo_use.operand_number - << " operand index " << hlo_use.operand_index.ToString() - << " size " << allocation_value.size() << " live range (" - << live_range_start_time << ", " << latest_prefetch_time - << ") from " - << (preferred_prefetch_time.has_value() - ? preferred_prefetch_time.value() - : -1) - << " to " - << (prefetch_override_info->value().prefetch_time.has_value() - ? absl::StrCat( - *prefetch_override_info->value().prefetch_time) - : "none"); - preferred_prefetch_time = prefetch_override_info->value().prefetch_time; - strict_timing = prefetch_override_info->value().strict_timing; - fail_on_unsatisfied_override = - prefetch_override_info->value().fail_on_unsatisfied_override; - is_prefetch_override = true; - } - - // Rarely, (e.g., when conditional true and false parameters are the - // same), definition time can be the time of the conditional and use - // time is the parameter use, which is less. - request.inclusive_start_time = std::min(definition_time, use_time); - request.latest_prefetch_time = latest_prefetch_time; - request.size = allocation_value_to_update.size(); - request.prefer_no_copy_alternate_mem_allocation = - prefer_no_copy_alternate_mem_allocation || - require_no_copy_alternate_mem_allocation; - request.allow_no_copy_alternate_mem_allocation = - allow_no_copy_alternate_mem_allocation; - request.allow_prefetch = allow_prefetch; - request.require_no_copy_alternate_mem_allocation = - require_no_copy_alternate_mem_allocation; - request.earliest_prefetch_time = earliest_prefetch_time; - request.preferred_prefetch_time = preferred_prefetch_time; - request.preferred_offset = preferred_offset; - request.use = &use; - request.allocation_value = &allocation_value; - request.all_use_times = all_use_times; - request.require_copy_allocation = require_copy_allocation; - request.required_copy_allocation_latest_time = - required_copy_allocation_latest_time; - request.required_copy_allocation_for = required_copy_allocation_for; - request.required_copy_for_slice = required_copy_for_slice; - request.allocation_value_to_update = &allocation_value_to_update; - request.strict_timing = strict_timing; - request.fail_on_unsatisfied_override = fail_on_unsatisfied_override; - request.is_prefetch_override = is_prefetch_override; - request.require_end_colored_in_default_memory = - require_end_colored_in_default_memory; - } - - if (shape_override.has_value()) { - request.shape_override = shape_override; - } - request.end_time = use_time; - request.only_extend_existing_allocation = only_extend_existing_allocation; - request.processed_allocation_values = processed_allocation_values; - return request; -} - -bool MsaAlgorithm::TryUpdateAllocationRequirementForUseAliases( - const AllocationValue& allocation_value, const AllocationValue::Use& use, - int64_t use_time) { - Allocation* aliased_allocation = - GetLiveAllocationAt(*allocation_value.allocation_sequence(), use_time); - VLOG(4) << "Aliased allocation at time " << use_time << ": " - << (aliased_allocation ? aliased_allocation->ToString() - : "couldn't find the aliased allocation"); - - if (!aliased_allocation) { - return true; - } - - // Before recording any aliased-use requirement, check every aliased - // position for a contradiction with an already recorded required - // assignment (or a buffer coloring). A contradiction means this use's - // allocation landed in a memory space that some aliased position cannot - // follow, e.g. among positions that must be colocated because an in-place - // instruction (such as a dynamic-update-slice) aliases them, this use's - // allocation landed in alternate memory while another aliased use already - // required the same position in default memory at the same time. That is - // a failed allocation, not a program invariant violation: report it so the - // caller can uncommit and retry (ultimately falling back to default - // memory) instead of CHECK-crashing inside AddRequiredAssignment. - AliasedOffset* aliased_offset = - aliased_allocation->memory_space() == MemorySpace::kAlternate - ? GetAliasedOffset(*aliased_allocation) - : nullptr; - for (const HloPosition& aliased_position : use.aliases) { - const HloValue* value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - aliased_position.instruction, aliased_position.index); - const int64_t time = - hlo_live_range_.instruction_schedule().at(aliased_position.instruction); - const RequiredMemoryAssignment new_requirement{ - aliased_allocation->memory_space(), time, aliased_offset, - RequiredMemoryAssignment::Source::kAliasedUse}; - const std::optional existing_requirement = - RequiredMemoryAssignmentAt(value, time); - if (existing_requirement.has_value() && - !new_requirement.memory_space_and_offset_equal(*existing_requirement)) { - VLOG(2) << "Aliased use requirement for " << value->ToShortString() - << " at time " << time << " (" << new_requirement.ToString() - << ") conflicts with the existing required assignment (" - << existing_requirement->ToString() - << "); failing this allocation so it can be retried."; - return false; - } - if (aliased_allocation->memory_space() == MemorySpace::kDefault && - IsPositionColoredInAlternateMemoryAtTime(value->defining_position(), - time)) { - VLOG(2) << "Aliased use requirement in default memory for " - << value->ToShortString() << " at time " << time - << " conflicts with an alternate memory coloring; failing this " - "allocation so it can be retried."; - return false; - } - if (aliased_allocation->memory_space() == MemorySpace::kAlternate && - IsPositionColoredInDefaultMemoryAtTime(value->defining_position(), - time)) { - VLOG(2) << "Aliased use requirement in alternate memory for " - << value->ToShortString() << " at time " << time - << " conflicts with a default memory coloring; failing this " - "allocation so it can be retried."; - return false; - } - } - - for (const HloPosition& aliased_position : use.aliases) { - AddAliasedRequiredAssignment(aliased_position.instruction, - aliased_position.index, aliased_allocation, - RequiredMemoryAssignment::Source::kAliasedUse); - } - return true; -} - -void MsaAlgorithm::MaybeCreateMirroredParentAllocationForWhileUse( - const AllocationValue& allocation_value, const AllocationValue::Use& use, - int64_t use_time, absl::Span allocation_values, - absl::flat_hash_map& - preferred_offset_for_computation) { - const HloUse& hlo_use = use.hlo_use; - - if (hlo_use.instruction->opcode() != HloOpcode::kWhile) { - return; - } - - Allocation* aliased_allocation = - GetLiveAllocationAt(*allocation_value.allocation_sequence(), use_time); - if (aliased_allocation->memory_space() != MemorySpace::kAlternate) { - return; - } - - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - if (options_.enable_while_redundant_eviction_elimination && - absl::c_find_if(allocation_value.value()->positions(), - [&hlo_use](const HloPosition& position) { - return position.instruction == hlo_use.instruction && - position.index == hlo_use.operand_index; - }) != allocation_value.value()->positions().end()) { - auto allocation_sequence = allocation_value.allocation_sequence(); - auto prev_allocation_in_default_mem_it = std::find_if( - allocation_sequence->rbegin(), allocation_sequence->rend(), - [&](const auto& allocation) { - return allocation->memory_space() == MemorySpace::kDefault && - allocation->defining_position() == allocation_value.position(); - }); - if (prev_allocation_in_default_mem_it != allocation_sequence->rend()) { - VLOG(3) << "Found a prev allocation in default mem for while use: " - << (*prev_allocation_in_default_mem_it)->ToString(); - // Safely match while loop parameters and outputs by exact tuple - // ShapeIndex using conditionals instead of rigid assertions, since - // unused return positions may not have active entries in - // allocation_values. - auto body_allocation_value_it = - absl::c_find_if(allocation_values, [&](const AllocationValue& value) { - return value.computation() == hlo_use.instruction->while_body() && - value.position().instruction->opcode() == - HloOpcode::kParameter && - value.position().index == hlo_use.operand_index; - }); - CHECK_NE(body_allocation_value_it, allocation_values.end()); - VLOG(3) << "Body allocation value: " - << body_allocation_value_it->ToShortString(); - int64_t body_parameter_time = instruction_schedule.at( - body_allocation_value_it->position().instruction); - body_allocation_value_it->mutable_allocation_sequence()->push_back( - std::make_unique( - **prev_allocation_in_default_mem_it, hlo_use.instruction, - body_allocation_value_it->position(), body_parameter_time)); - VLOG(3) << "Created: " - << body_allocation_value_it->allocation_sequence() - ->back() - ->ToString(); - - auto after_while_allocation_value_it = - absl::c_find_if(allocation_values, [&](const AllocationValue& value) { - return value.position().instruction == hlo_use.instruction && - value.position().index == hlo_use.operand_index; - }); - CHECK_NE(after_while_allocation_value_it, allocation_values.end()); - VLOG(3) << "After while allocation value: " - << after_while_allocation_value_it->ToShortString(); - int64_t while_time = instruction_schedule.at(hlo_use.instruction); - after_while_allocation_value_it->mutable_allocation_sequence()->push_back( - std::make_unique( - **prev_allocation_in_default_mem_it, while_time)); - VLOG(3) << "Created: " - << after_while_allocation_value_it->allocation_sequence() - ->back() - ->ToString(); - } - } - // Special case for while loops since the root offset must agree with - // other offsets: remember the preferred offset for the while loop body. - AliasedOffset* offset = GetAliasedOffset(*aliased_allocation); - if (IsAsyncPipelinedWhileLoop(hlo_use.instruction)) { - SynchronizeAliasedWhileLoopOffsets(hlo_use, *aliased_allocation, offset); - } else { - preferred_offset_for_computation[hlo_use.instruction->while_body()] = - offset; - } -} - -void MsaAlgorithm::SynchronizeAliasedWhileLoopOffsets( - const HloUse& hlo_use, const Allocation& aliased_allocation, - AliasedOffset* offset) { - HloComputation* while_body = hlo_use.instruction->while_body(); - pipelined_while_preferred_offset_for_computation_[while_body] - [hlo_use.operand_index] = - offset; - - // Enforce strict aliasing synchronization across the loop parameter, body - // root, and external while instruction by recording the canonical - // AliasedOffset against all associated positions. In XLA HLO, a while loop - // body computation always takes a single tuple parameter (index 0) - // representing the entire loop state. - static constexpr int64_t kWhileLoopStateParameterNumber = 0; - HloInstruction* while_parameter = - while_body->parameter_instruction(kWhileLoopStateParameterNumber); - AddAliasedRequiredAssignment(while_parameter, hlo_use.operand_index, - &aliased_allocation, - RequiredMemoryAssignment::Source::kAliasedUse); - RecordAliasedOffsetForAsyncPipelinedWhileLoop( - {while_parameter, hlo_use.operand_index}, offset); - - AddAliasedRequiredAssignment(while_body->root_instruction(), - hlo_use.operand_index, &aliased_allocation, - RequiredMemoryAssignment::Source::kAliasedUse); - RecordAliasedOffsetForAsyncPipelinedWhileLoop( - {while_body->root_instruction(), hlo_use.operand_index}, offset); - - AddAliasedRequiredAssignment(hlo_use.instruction, hlo_use.operand_index, - &aliased_allocation, - RequiredMemoryAssignment::Source::kAliasedUse); - RecordAliasedOffsetForAsyncPipelinedWhileLoop( - {hlo_use.instruction, hlo_use.operand_index}, offset); -} - -bool MsaAlgorithm::ShouldBeMirroredByNestedConditionals( - const AllocationValue& allocation_value, - const AllocationValue::Use& current_use, - const AllocationValue::Use* previous_use) const { - // If all of the following conditions are met, we need mirrored allocations - // for the aliasing, nested allocation values in the time range - // [time(`previous_use`), time(`current_use`)]: - // 1. The previous use is a conditional and the current use is strictly after - // the conditional. - // 2. The previous allocation in the AllocationSequence is in the alternate - // memory. - // 3. The previous allocation serves the previous use and the current use. - // 4. The previous allocation extends throughout the conditional live range. - - // Check conditions 1 and 2. - const AllocationSequence* allocation_sequence = - allocation_value.allocation_sequence(); - if (previous_use == nullptr || - previous_use->hlo_use.instruction->opcode() != HloOpcode::kConditional || - current_use.hlo_use.instruction == previous_use->hlo_use.instruction || - allocation_sequence->empty() || - allocation_sequence->back()->memory_space() != MemorySpace::kAlternate) { - return false; - } - // Check condition 3. - const Allocation* last_allocation = allocation_sequence->back().get(); - if (!absl::c_linear_search(last_allocation->uses(), previous_use->hlo_use)) { - return false; - } - // Check condition 4. - int64_t conditional_start_time = GetCorrectedUseTime(previous_use->hlo_use); - int64_t conditional_end_time = hlo_live_range_.instruction_schedule().at( - previous_use->hlo_use.instruction); - int64_t current_use_time = GetCorrectedUseTime(current_use.hlo_use); - bool last_allocation_covers_conditional_live_range = - last_allocation->earliest_available_time() <= conditional_start_time && - conditional_end_time <= current_use_time && - current_use_time <= last_allocation->end_time(); - return last_allocation_covers_conditional_live_range; -} - -bool MsaAlgorithm::ShouldBeMirrored( - const AllocationValue& allocation_value, - const AllocationValue::Use& current_use, - const AllocationValue::Use* previous_use) const { - return ShouldBeMirroredByNestedConditionals(allocation_value, current_use, - previous_use); -} - -void MsaAlgorithm::CreateMirroredAllocationsForNestedConditionals( - AllocationValue& allocation_value, const AllocationValue::Use& current_use, - const AllocationValue::Use* previous_use, - absl::Span allocation_values, - absl::flat_hash_set& - already_processed_allocation_values_inside_a_conditional) { - CHECK_NE(previous_use, nullptr); - const Allocation* last_allocation = - allocation_value.allocation_sequence()->back().get(); - int64_t conditional_start_time = GetCorrectedUseTime(previous_use->hlo_use); - int64_t conditional_end_time = hlo_live_range_.instruction_schedule().at( - previous_use->hlo_use.instruction); - - const HloBuffer* buffer_in_alt_mem = &alias_analysis_.GetUniqueBufferAt( - allocation_value.value()->defining_instruction(), - allocation_value.value()->defining_position().index); - - // Find allocations of the same buffer that lie within the conditional live - // range, create mirrored allocations in alternate memory, and add them to the - // set of processed allocation values throughout conditionals, to avoid - // re-processing. - for (AllocationValue& allocation_val : allocation_values) { - // 1. Find allocation values whose positions lie inside the live - // range of the conditional. - int64_t position_time = hlo_live_range_.instruction_schedule().at( - allocation_val.position().instruction); - int64_t last_use_time = position_time; - for (const AllocationValue::Use& use : allocation_val.uses()) { - last_use_time = std::max( - last_use_time, - hlo_live_range_.instruction_schedule().at(use.hlo_use.instruction)); - } - if (position_time < conditional_start_time || - position_time >= conditional_end_time) { - continue; - } - CHECK(conditional_start_time <= last_use_time && - last_use_time <= conditional_end_time) - << "last_use_time: " << last_use_time - << ", conditional_start_time: " << conditional_start_time - << ", conditional_end_time: " << conditional_end_time - << "\n allocation_val: " << allocation_val.ToString() - << "\n allocation_value: " << allocation_value.ToString() - << "\n previous_use: " << previous_use->hlo_use.ToString() - << "\n current use: " << current_use.hlo_use.ToString(); - - // 2. Add the allocation value to the set of allocation values that are - // live throughout a conditional and processed, so we do not process them - // again, later. - already_processed_allocation_values_inside_a_conditional.insert( - &allocation_val); - - // 3. If the allocation value is an async conversion candidate, create an - // allocation in default memory to avoid a copy from alternate memory - // to alternate memory. Additionally, we check if the source buffer of - // the async conversion candidate is the same as the buffer in alternate - // memory. - HloInstruction* allocation_value_instruction = - allocation_val.value()->defining_instruction(); - if (IsAsyncConversionCandidate(allocation_value_instruction)) { - const HloBuffer* source_buffer = &alias_analysis_.GetUniqueBufferAt( - allocation_value_instruction->mutable_operand(0), {}); - CHECK_EQ(source_buffer, buffer_in_alt_mem) - << "Source buffer mismatch for async conversion candidate: " - << allocation_value_instruction->ToString() - << " \n with buffer: " << buffer_in_alt_mem->ToString(); - VLOG(3) - << "Skipping async conversion candidate because the source buffer is " - "in alternate memory and a successful async conversion would " - "result in a copy from alternate memory to alternate memory." - << allocation_val.ToString(); - allocation_val.mutable_allocation_sequence()->push_back( - std::make_unique(allocation_val.position(), - MemorySpace::kDefault, kDummyChunk, - position_time, last_use_time)); - for (const AllocationValue::Use& use : allocation_val.uses()) { - allocation_val.mutable_allocation_sequence()->back()->AddUse( - use.hlo_use); - } - continue; - } - - VLOG(3) << "Conditional operand: " << previous_use->hlo_use.ToString() - << " is in alternate memory throughout the conditional, adding a" - " mirrored allocation for pre-allocated allocation_value: " - << allocation_val.ToString(); - - // 4. We create a mirrored allocation in the alternate memory for the - // allocation value. - allocation_val.mutable_allocation_sequence()->push_back( - std::make_unique(allocation_val.position(), - *last_allocation, position_time, - last_use_time)); - for (const AllocationValue::Use& use : allocation_val.uses()) { - allocation_val.mutable_allocation_sequence()->back()->AddUse(use.hlo_use); - } - MaybeCreateOrAddToAliasedOffset( - *allocation_val.mutable_allocation_sequence()->back(), - /*aliased_offset=*/nullptr); - } -} - -void MsaAlgorithm::CreateMirroredAllocations( - AllocationValue& allocation_value, const AllocationValue::Use& current_use, - const AllocationValue::Use* previous_use, - absl::Span allocation_values, - absl::flat_hash_set& - already_processed_allocation_values_inside_a_conditional) { - CreateMirroredAllocationsForNestedConditionals( - allocation_value, current_use, previous_use, allocation_values, - already_processed_allocation_values_inside_a_conditional); -} - -bool MsaAlgorithm::IsEvictionRequiredForPreviousUseAtConditional( - AllocationValue& allocation_value, const AllocationValue::Use& use, - const AllocationValue::Use* previous_use, - absl::Span allocation_values) { - if (previous_use == nullptr || - previous_use->hlo_use.instruction->opcode() != HloOpcode::kConditional || - use.hlo_use.instruction == previous_use->hlo_use.instruction) { - return false; - } - - int64_t conditional_end_time = hlo_live_range_.instruction_schedule().at( - previous_use->hlo_use.instruction); - int64_t current_use_time = GetCorrectedUseTime(use.hlo_use); - - if (conditional_end_time > current_use_time) { - return false; - } - - const HloBuffer* buffer_in_alt_mem = &alias_analysis_.GetUniqueBufferAt( - allocation_value.value()->defining_instruction(), - allocation_value.value()->defining_position().index); - int64_t conditional_start_time = GetCorrectedUseTime(previous_use->hlo_use); - - for (AllocationValue& allocation_val : allocation_values) { - int64_t position_time = hlo_live_range_.instruction_schedule().at( - allocation_val.position().instruction); - int64_t last_use_time = position_time; - for (const AllocationValue::Use& use : allocation_val.uses()) { - last_use_time = std::max( - last_use_time, - hlo_live_range_.instruction_schedule().at(use.hlo_use.instruction)); - } - if (position_time < conditional_start_time || - position_time >= conditional_end_time) { - continue; - } - CHECK(conditional_start_time <= last_use_time && - last_use_time <= conditional_end_time) - << "last_use_time: " << last_use_time - << ", conditional_start_time: " << conditional_start_time - << ", conditional_end_time: " << conditional_end_time - << "\n allocation_val: " << allocation_val.ToString() - << "\n allocation_value: " << allocation_value.ToString() - << "\n previous_use: " << previous_use->hlo_use.ToString() - << "\n current use: " << use.hlo_use.ToString(); - - if (IsAsyncConversionCandidate( - allocation_val.value()->defining_instruction())) { - // Check the jointly processed allocation values within the conditional - // live range. Allow processing without eviction, if and only if, all - // the jointly processed async conversion candidates have the source - // buffer as the buffer in alternate memory. We will handle this case by - // forcing the async conversion candidate to be in default memory. We - // force all other kinds of jointly processed allocation values that - // come from async conversion candidates to require an eviction because - // async copying to and from the same memory space is not allowed. - // Trying to satisfy this restriction might impose infeasible - // constraints on how other buffers must get allocated. - const HloBuffer* source_buffer_of_copy = - &alias_analysis_.GetUniqueBufferAt( - allocation_val.value()->defining_instruction()->mutable_operand( - 0), - {}); - if (source_buffer_of_copy == buffer_in_alt_mem) { - VLOG(3) - << "Source buffer is the same as the buffer in alternate memory: " - << allocation_val.ToString(); - continue; - } - return true; - } - - // Require eviction if the set of allocation values jointly processed - // within the conditional live range span across different buffers. - const HloBuffer* destination_buffer = &alias_analysis_.GetUniqueBufferAt( - allocation_val.position().instruction, allocation_val.position().index); - if (buffer_in_alt_mem != destination_buffer) { - CHECK(!IsAsyncConversionCandidate( - allocation_val.value()->defining_instruction())); - return true; - } - - for (const AllocationValue::Use& use : allocation_val.uses()) { - if (!IsUseAllowedInAlternateMemory(use.hlo_use) || - !IsConditionalUseBeneficialInAlternateMemory(allocation_val, - use.hlo_use)) { - VLOG(3) << "Use not allowed in alternate memory: " - << use.hlo_use.ToString() - << " requiring eviction for allocation_val: " - << allocation_val.ToString() - << " at conditional use: " << previous_use->hlo_use.ToString(); - return true; - }; - std::optional required_assignment_for_use = - RequiredAssignmentForUse(use); - if (required_assignment_for_use.has_value() && - required_assignment_for_use.value().memory_space == - MemorySpace::kDefault) { - VLOG(3) << "Required assignment: " - << required_assignment_for_use.value().ToString() - << " for use: " << use.hlo_use.ToString() - << " in default memory, requiring eviction for allocation_val: " - << allocation_val.ToString() - << " at conditional use: " << previous_use->hlo_use.ToString(); - return true; - } - std::optional aliased_required_assignment = - AliasedRequiredAssignmentForUse(use); - if (aliased_required_assignment.has_value() && - aliased_required_assignment.value().memory_space == - MemorySpace::kDefault) { - VLOG(3) << "Aliased required assignment: " - << aliased_required_assignment.value().ToString() - << " for use: " << use.hlo_use.ToString() - << " in default memory, requiring eviction for allocation_val: " - << allocation_val.ToString() - << " at conditional use: " << previous_use->hlo_use.ToString(); - return true; - } - std::optional - required_assignment_at_defining_instruction = - RequiredMemoryAssignmentAt(allocation_val.value(), position_time); - if (required_assignment_at_defining_instruction.has_value() && - required_assignment_at_defining_instruction.value().memory_space == - MemorySpace::kDefault) { - VLOG(3) - << "Required assignment at defining instruction: " - << required_assignment_at_defining_instruction.value().ToString() - << " for allocation_val: " << allocation_val.ToString() - << " requiring eviction at conditional use: " - << previous_use->hlo_use.ToString(); - return true; - } - } - } - return false; -} - -bool operator<(const AsynchronousCopy& a, const AsynchronousCopy& b) { - return a.AsTuple() < b.AsTuple(); -} - -bool operator==(const AsynchronousCopy& a, const AsynchronousCopy& b) { - return a.AsTuple() == b.AsTuple(); -} - -bool operator!=(const AsynchronousCopy& a, const AsynchronousCopy& b) { - return a.AsTuple() != b.AsTuple(); -} - -void AsynchronousCopyOrdering::AddCopy(const AsynchronousCopy& copy) { - auto it = ranges_.find({copy.exclusive_start_time, copy.end_time}); - if (it != ranges_.end()) { - CHECK_EQ(it->first.exclusive_start_time, copy.exclusive_start_time); - CHECK(it->second.insert(copy).second); - } else { - ranges_[{copy.exclusive_start_time, copy.end_time}] = {copy}; - } -} - -void AsynchronousCopyOrdering::RemoveCopy(const AsynchronousCopy& copy) { - auto copy_it = ranges_.find({copy.exclusive_start_time, copy.end_time}); - CHECK(copy_it != ranges_.end()); - CHECK_EQ(copy_it->first.exclusive_start_time, copy.exclusive_start_time); - CHECK_EQ(copy_it->second.erase(copy), 1); - if (copy_it->second.empty()) { - ranges_.erase(copy_it); - } -} - -bool AsynchronousCopyOrdering::ViolatesOrdering(int64_t exclusive_start_time, - int64_t end_time) const { - // We allow identical start and end times. It is enough to check for just the - // start time in case we find a match in ranges_ because the found value will - // either be identical to {start_time, estimated_end_time} (and this doesn't - // violate) or its start_time will be smaller and estimated_end_time will be - // larger (this violates). - auto copy_it = ranges_.find({exclusive_start_time, end_time}); - if (copy_it != ranges_.end() && - copy_it->first.exclusive_start_time != exclusive_start_time) { - VLOG(4) << "Violates ordering: (" << exclusive_start_time << ", " - << end_time << ") and (" << copy_it->first.exclusive_start_time - << ", " << copy_it->first.end_time << ")"; - return true; - } - return false; -} - -bool AsynchronousCopyResource::ConsumeResource( - int64_t exclusive_start_time, int64_t end_time, int64_t resource, - std::vector>* delay_changes, - int64_t resource_to_free) { - // Cache the pointers to the arrays to avoid the overhead of `operator[]` - // size checks in hardened libc++. - // - // NOTE: Do not modify the vectors `initial_resources_` or `delay_` in this - // function, otherwise the pointers will become dangling. - int64_t* initial_resources_scaled_ptr = initial_resources_scaled_.data(); - int64_t* delay_ptr = delay_.data(); - - auto current_copy = async_copies_.end(); - // In order to propagate the resource to the next scheduled copy, we iterate - // over the copies in start time order until we either find enough free - // resource (and return true), or find out that we don't have enough free - // resource (and return false). - while (true) { - // resource is modified below. We save its initial value for logging below. - const int64_t amount_requested = resource; - - VLOG(3) << "Consume resource: start time_exclusive = " - << exclusive_start_time << ", end time = " << end_time - << ", resource = " << resource << ", delay = " - << delay_[ExclusiveToInclusiveStartTime(exclusive_start_time)] - << ", free = " << resource_to_free; - VLOG(5) << "Available resources: " - << VectorToString( - GetCurrentResources(), /*include_indices=*/true, - ExclusiveToInclusiveStartTime(exclusive_start_time), - end_time); - - // Nothing to do if we're not adding or removing any resources. - if (resource == 0 && resource_to_free == 0) { - return true; - } - - // For the async copy we're adding, check the delay_ array to see how much - // this copy would have to be delayed because of an earlier copy that wasn't - // finished when this copy starts. - if (current_copy == async_copies_.end()) { - resource += - delay_ptr[ExclusiveToInclusiveStartTime(exclusive_start_time)]; - } - - // Find the copy that is right after this one. If there are leftover - // resources by the time the next copy starts, the next copy will be pushed - // further later in time. - auto next_copy = async_copies_.end(); - if (current_copy != async_copies_.end()) { - next_copy = std::next(current_copy); - } else { - auto async_copy_time_it = - async_copy_time_map_.upper_bound(exclusive_start_time); - if (async_copy_time_it != async_copy_time_map_.end()) { - next_copy = async_copy_time_it->second; - } - } - - // Check if this copy will push the next copy later in time (or if removing - // the resource, check if the removal of this copy move the next copy - // earlier in time). - std::optional delay_for_next_copy = std::nullopt; - int64_t resource_freed = 0; - for (int64_t time = ExclusiveToInclusiveStartTime(exclusive_start_time); - time < end_time && resource != 0; ++time) { - int64_t initial_resource_scaled = initial_resources_scaled_ptr[time]; - // Iterate over the logical times that this copy spans. Note that the - // start and end time ranges are exclusive. - int64_t used_resource = std::min(resource, initial_resource_scaled); - if (next_copy != async_copies_.end() && - next_copy->exclusive_start_time == - InclusiveToExclusiveStartTime(time)) { - // This is the time where the next copy begins. If the resource is - // non-zero at this point, the copy didn't finish by the time the next - // copy started, so the next copy would need to be pushed later in time. - delay_for_next_copy = resource; - resource_to_free -= resource_freed; - } - if (!delay_for_next_copy.has_value()) { - // Update the delay_ vector and resource_freed variable with the amount - // that was freed when removing the copy. - int64_t old_delay = delay_ptr[time]; - int64_t old_resource = - std::max(0, initial_resource_scaled - old_delay); - int64_t new_delay = std::max(0, resource - resource_to_free); - int64_t new_resource = - std::max(0, initial_resource_scaled - new_delay); - resource_freed += std::max(0, new_resource - old_resource); - delay_ptr[time] = new_delay; - if (delay_changes) { - delay_changes->emplace_back(time, old_delay); - } - } - // Update the resource with the used amount in this logical time. - resource -= used_resource; - } - - // If resource isn't satisfied by the end, we didn't have enough - // resources. - if (resource > 0) { - VLOG(3) << "Doesn't have enough resource; requested resource = " - << amount_requested << "; leftover resources = " << resource; - return false; - } - - if (!delay_for_next_copy.has_value()) { - return true; - } - // If this copy overlapped with another one, we run for another iteration - // with the next copy with the amount of resource that needs to be added or - // removed. - exclusive_start_time = next_copy->exclusive_start_time; - end_time = next_copy->end_time; - resource = - *delay_for_next_copy + GetScaledIntegerResource(next_copy->resource); - current_copy = next_copy; - } -} - -void AsynchronousCopyResource::AddCopy(const AsynchronousCopy& copy) { - CHECK(ConsumeResource(copy.exclusive_start_time, copy.end_time, - GetScaledIntegerResource(copy.resource))); - - // Find the iterator for the copy that would be right after this copy and put - // this copy right before it in async_copies_. - auto async_copy_time_it = - async_copy_time_map_.upper_bound(copy.exclusive_start_time); - auto insertion_it = (async_copy_time_it == async_copy_time_map_.end()) - ? async_copies_.end() - : async_copy_time_it->second; - auto inserted_it = async_copies_.insert(insertion_it, copy); - // If this copy is the first copy we have seen with the start time, add the - // inserted iterator into async_copy_time_map_ for fast lookups. Note that - // async_copy_time_map_ always points to the very first copy with the same - // start index. If there are multiple asynchronous copies that have the same - // start time, the memory space assignment algorithm schedules them in the - // same order that AddCopy was called. - if (async_copy_time_map_.find(copy.exclusive_start_time) == - async_copy_time_map_.end()) { - async_copy_time_map_[copy.exclusive_start_time] = inserted_it; - } -} - -void AsynchronousCopyResource::RemoveCopy(const AsynchronousCopy& copy) { - // The ConsumeResource method can only correctly remove the last copy that - // starts at a given start time. So if the copy that is requested to be - // removed is not the last copy for this start time, we need to temporarily - // remove later copies that has the same start time and then add them back one - // by one. To do this, we first find the iterator that points to the earliest - // copy after this start time. We then decrement this iterator and temporarily - // remove the copies until we find the copy we actually want to remove. After - // we remove the copy that we actually want to remove, we add back the - // temporarily removed copies one by one in the same order. - auto async_copy_time_it = - async_copy_time_map_.upper_bound(copy.exclusive_start_time); - auto copy_it = (async_copy_time_it == async_copy_time_map_.end()) - ? async_copies_.end() - : async_copy_time_it->second; - CHECK(copy_it != async_copies_.begin()); - --copy_it; - - std::list copies_to_add_back; - auto prev_copy_it = copy_it; - for (; *copy_it != copy; copy_it = prev_copy_it) { - CHECK(copy_it != async_copies_.begin()); - CHECK_EQ(copy_it->exclusive_start_time, copy.exclusive_start_time); - copies_to_add_back.push_front(*copy_it); - VLOG(4) << "RemoveCopy found a copy to temporarily remove and add back: " - << copy_it->exclusive_start_time << " " << copy_it->end_time << " " - << copy_it->resource; - prev_copy_it = std::prev(copy_it); - RemoveCopy(copy_it); - } - CHECK(*copy_it == copy); - RemoveCopy(copy_it); - - for (const AsynchronousCopy& copy_to_add_back : copies_to_add_back) { - AddCopy(copy_to_add_back); - } -} - -void AsynchronousCopyResource::RemoveCopy( - std::list::iterator& copy_it) { - // This method works only for the latest copy for the given start time. - CHECK(std::next(copy_it) == async_copies_.end() || - std::next(copy_it)->exclusive_start_time > - copy_it->exclusive_start_time); - CHECK(ConsumeResource( - copy_it->exclusive_start_time, copy_it->end_time, - /*resource=*/0, - /*delay_changes=*/nullptr, - /*resource_to_free=*/GetScaledIntegerResource(copy_it->resource))); - // If the copy to be removed is the value pointed by async_copy_time_map_, we - // make the next copy with the same start time to be pointed by - // async_copy_time_map_. If there are no such copies, we remove the key for - // this copy start time. - int64_t exclusive_start_time = copy_it->exclusive_start_time; - auto async_copy_time_it = async_copy_time_map_.find(exclusive_start_time); - if (copy_it == async_copy_time_it->second) { - if (std::next(copy_it) != async_copies_.end() && - std::next(copy_it)->exclusive_start_time == exclusive_start_time) { - async_copy_time_it->second = std::next(copy_it); - } else { - async_copy_time_map_.erase(async_copy_time_it); - } - } - async_copies_.erase(copy_it); -} - -bool AsynchronousCopyResource::HasEnoughResource(int64_t exclusive_start_time, - int64_t end_time, - float resource) { - std::vector> delay_changes; - delay_changes.reserve(delay_.size()); - bool result = - ConsumeResource(exclusive_start_time, end_time, - GetScaledIntegerResource(resource), &delay_changes); - // Apply the delay changes in reverse order. This ensures that the original - // value of each delay is restored. - if (!delay_changes.empty()) { - for (int64_t i = delay_changes.size() - 1; i >= 0; --i) { - const auto& [time, delay] = delay_changes[i]; - delay_[time] = delay; - } - } - return result; -} - -bool AsynchronousCopyResource::HasEnoughResourceMultiCheck( - const std::vector& specs) { - delay_changes_.resize(0); - delay_changes_.reserve(delay_.size()); - bool result = absl::c_all_of(specs, [&](const ResourceSpec& spec) { - return ConsumeResource(spec.exclusive_start_time, spec.end_time, - GetScaledIntegerResource(spec.resource), - &delay_changes_); - }); - // Apply the delay changes in reverse order. This ensures that the original - // value of each delay is restored. - if (!delay_changes_.empty()) { - for (int64_t i = delay_changes_.size() - 1; i >= 0; --i) { - const auto& [time, delay] = delay_changes_[i]; - delay_[time] = delay; - } - } - return result; -} - -namespace { - -// A convenience struct for use in the implementation of -// AsynchronousCopyResource::Dump(). -struct CopyResourceDumpData { - float initial_resource; - float delay; - float available; - std::vector overlapping_copies; -}; - -} // namespace - -std::string AsynchronousCopyResource::Dump( - int64_t start_time, int64_t end_time, - MemorySpace memory_space_filter) const { - std::vector available = GetCurrentResources(); - std::vector time_dump_data; - for (int i = start_time; i < end_time; ++i) { - time_dump_data.push_back({ - initial_resources_[i], - GetDescaledFloatResource(delay_[i]), - available[i], - /*overlapping_copies=*/{}, - }); - } - - std::vector lines; - lines.push_back(absl::StrCat("AsynchronousCopyResource::Dump(start_time: ", - start_time, ", end_time: ", end_time, ")")); - for (const AsynchronousCopy& copy : async_copies_) { - if (copy.destination != memory_space_filter) { - continue; - } - int64_t overlap_start = std::max(start_time, copy.exclusive_start_time); - int64_t overlap_end = std::min(end_time, copy.end_time); - if (overlap_start < overlap_end) { - lines.push_back(absl::StrCat( - "copy(id: ", copy.id, - ", exclusive_start: ", copy.exclusive_start_time, - ", end: ", copy.end_time, ", resource: ", copy.resource, ")")); - } - for (int i = overlap_start; i < overlap_end; ++i) { - time_dump_data[i - start_time].overlapping_copies.push_back(copy.id); - } - } - - std::vector col_sizes; - std::vector> rows; - rows.push_back({"time", "initial", "delay", "avail", "overlapping copies"}); - for (absl::string_view col : rows.front()) { - col_sizes.push_back(col.size()); - } - for (int i = 0; i < time_dump_data.size(); ++i) { - rows.push_back({absl::StrCat(i + start_time), - absl::StrCat(time_dump_data[i].initial_resource), - absl::StrCat(time_dump_data[i].delay), - absl::StrCat(time_dump_data[i].available), - absl::StrJoin(time_dump_data[i].overlapping_copies, ",")}); - for (int j = 0; j < rows.back().size(); ++j) { - col_sizes[j] = std::max(col_sizes[j], rows.back()[j].size()); - } - } - for (const std::vector& row : rows) { - std::string line; - std::string sep; - for (int i = 0; i < col_sizes.size(); ++i) { - absl::StrAppend(&line, sep, row[i]); - sep = std::string(col_sizes[i] + 2 - row[i].size(), ' '); - } - lines.push_back(line); - } - - return absl::StrJoin(lines, "\n"); -} - -AliasedOffset* MsaAlgorithm::GetAliasedOffset(const Allocation& allocation) { - if (allocation.is_mirrored_allocation()) { - const MirroredAllocation* mirrored_allocation = - dynamic_cast(&allocation); - CHECK(mirrored_allocation != nullptr); - return GetAliasedOffset(mirrored_allocation->original_allocation()); - } - auto aliased_offset_it = aliased_offset_map_.find(&allocation); - CHECK(aliased_offset_it != aliased_offset_map_.end()) - << "Allocation not found in aliased offset map: " - << allocation.ToString(); - return aliased_offset_it->second; -} - -void MsaAlgorithm::MaybeCreateOrAddToAliasedOffset( - const Allocation& allocation, AliasedOffset* aliased_offset) { - CHECK(allocation.memory_space() == MemorySpace::kAlternate) - << "Allocation is not in the alternate memory: " << allocation.ToString(); - if (allocation.is_mirrored_allocation()) { - return; - } - CHECK(!aliased_offset_map_.contains(&allocation)) - << "Allocation already has an aliased offset: " << allocation.ToString(); - if (!aliased_offset) { - aliased_offsets_.push_back({allocation.chunk().offset}); - aliased_offset = &aliased_offsets_.back(); - } - CHECK_EQ(allocation.chunk().offset, aliased_offset->offset); - CHECK(aliased_offset->allocations.insert(&allocation).second); - aliased_offset_map_[&allocation] = aliased_offset; -} - -void MsaAlgorithm::RecordAliasedOffsetForAsyncPipelinedWhileLoop( - const HloPosition& position, AliasedOffset* aliased_offset) { - if (aliased_offset != nullptr) { - const HloBuffer& hlo_buffer = - alias_analysis_.GetUniqueBufferAt(position.instruction, position.index); - pipelined_while_buffer_id_to_aliased_offset_[hlo_buffer.id()] = - aliased_offset; - } -} - -/*static*/ Allocation* MsaAlgorithm::GetLiveAllocationAt( - const AllocationSequence& allocations, int64_t time) { - for (auto allocation_it = allocations.rbegin(); - allocation_it != allocations.rend(); ++allocation_it) { - // The use case of GetLiveAllocationAt is to find the allocation that - // corresponds to the full buffer. Window prefetched allocations allocates - // only partial buffers, so we want to skip them. - if ((*allocation_it)->start_time() <= time && - (*allocation_it)->end_time() >= time && - !(*allocation_it)->is_window_prefetched_allocation()) { - return allocation_it->get(); - } - } - return nullptr; -} - -namespace { - -void SetDefaultMemorySpace(const HloValue* value, const Options& options) { - for (auto& position : value->positions()) { - Shape* shape = ShapeUtil::GetMutableSubshape( - position.instruction->mutable_shape(), position.index); - if (!shape->has_layout() || - shape->layout().memory_space() != options.alternate_memory_space) { - continue; - } - shape->mutable_layout()->set_memory_space(options.default_memory_space); - } - HloModule* module = value->defining_instruction()->GetModule(); - module->mutable_config().SetComputationLayoutIfExists( - module->entry_computation()->ComputeProgramShape()); -} - -} // namespace - -void MsaAlgorithm::AllocateCrossProgramPrefetchBuffer( - HloModule* module, const MsaBufferInterval& prefetch_candidate) { - Chunk chunk_candidate = FindChunkCandidate(prefetch_candidate); - if (chunk_candidate.chunk_end() > available_heap_size()) { - VLOG(3) << "Could not allocate preferred memory for cross program prefetch"; - return; - } - - const HloValue* buffer = prefetch_candidate.buffer; - int64_t parameter = buffer->instruction()->parameter_number(); - int cross_program_prefetch_index = module->CrossProgramPrefetches().size(); - SetDefaultMemorySpace(buffer, options_); - module->AddCrossProgramPrefetch(parameter, buffer->index()); - - AllocationSequence allocations; - allocations.push_back(std::make_unique( - buffer->defining_position(), MemorySpace::kDefault, kDummyChunk, - prefetch_candidate.start, prefetch_candidate.end)); - - // Find the earliest use. - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - auto uses = FindCrossProgramPrefetchUses(buffer->GetUses(), alias_analysis_, - alias_info_); - CHECK_GE(uses.size(), 1); - auto use_schedule_compare = [&](const HloUse& lhs, const HloUse& rhs) { - return instruction_schedule.at(lhs.instruction) < - instruction_schedule.at(rhs.instruction); - }; - auto first_use = absl::c_min_element(uses, use_schedule_compare); - int64_t latest_prefetch_time = - instruction_schedule.at(first_use->instruction); - - // Find the last use time. Note, view uses may extend the last use time. - int64_t last_use_time = std::numeric_limits::min(); - for (const HloUse& use : uses) { - last_use_time = std::max(last_use_time, GetExtendedUseTimeIfUseIsView(use)); - } - for (const HloValue* colocation : prefetch_candidate.colocations) { - for (const HloUse& use : colocation->GetUses()) { - last_use_time = - std::max(last_use_time, GetExtendedUseTimeIfUseIsView(use)); - } - } - - int64_t end_of_program_prefetch_end_time = instruction_schedule.size(); - int64_t end_of_program_inclusive_prefetch_start_time = - options_.prefetch_interval_picker->PreferredPrefetchStartTime( - buffer->defining_position().shape(), last_use_time, - end_of_program_prefetch_end_time, end_of_program_prefetch_end_time); - VLOG(2) << "last use time = " << last_use_time - << ", end-of-program inclusive prefetch start time = " - << end_of_program_inclusive_prefetch_start_time; - float total_execution_time = - options_.prefetch_interval_picker->GetLogicalIntervalElapsed( - 0, instruction_schedule.size()); - float buffer_occupied_time = - options_.prefetch_interval_picker->GetLogicalIntervalElapsed( - end_of_program_inclusive_prefetch_start_time, - end_of_program_prefetch_end_time); - if (options_.cost_analysis) { - buffer_occupied_time = std::max( - buffer_occupied_time, options_.cost_analysis->GetAsyncCopyElapsed( - options_.cost_analysis->GetShapeSizeBytes( - buffer->defining_position().shape()))); - } - buffer_occupied_time += - options_.prefetch_interval_picker->GetLogicalIntervalElapsed( - 0, last_use_time); - float buffer_occupied_ratio = buffer_occupied_time / total_execution_time; - VLOG(2) << "Total execution time = " << total_execution_time - << ", buffer occupied time = " << buffer_occupied_time - << ", buffer occupied ratio = " << buffer_occupied_ratio; - // Freeing buffer only makes sense if the buffer will be free for a - // substantial time. Only perform this optimization if the ratio is below the - // limit, and if the memory pressure is above the alternate memory size. - bool free_buffer = - (options_.enable_cross_program_prefetch_freeing && - memory_pressure_ > options_.max_size_in_bytes && - buffer_occupied_ratio < kCrossProgramPrefetchOccupyFreeingLimit && - end_of_program_inclusive_prefetch_start_time > last_use_time && - end_of_program_inclusive_prefetch_start_time < - end_of_program_prefetch_end_time); - int64_t cross_program_prefetch_end_time = - free_buffer ? last_use_time : prefetch_candidate.end; - - AddAsyncCopyOrOtherMemOp( - *allocations.back(), MemorySpace::kAlternate, chunk_candidate, - /*exclusive_start_time=*/ - InclusiveToExclusiveStartTime(prefetch_candidate.start), - cross_program_prefetch_end_time, latest_prefetch_time, &allocations, - /*aliased_offset=*/nullptr, - /*resource=*/0.0, cross_program_prefetch_index); - - absl::c_for_each(uses, [&](auto& use) { allocations.back()->AddUse(use); }); - AliasedOffset* cross_program_prefetch_offset = - GetAliasedOffset(*allocations.back()); - - if (free_buffer) { - VLOG(2) << "Adding an end-of-program prefetch for freed " - "cross-program-prefetched buffer."; - AddAsyncCopyOrOtherMemOp( - *allocations.front(), MemorySpace::kAlternate, chunk_candidate, - /*exclusive_start_time=*/ - InclusiveToExclusiveStartTime( - end_of_program_inclusive_prefetch_start_time), - end_of_program_prefetch_end_time, end_of_program_prefetch_end_time, - &allocations, cross_program_prefetch_offset, - /*resource=*/0.0); - CHECK_EQ(cross_program_prefetch_offset->offset, - allocations.back()->chunk().offset); - } - - const int allocations_initial_size = allocations_->size(); - for (auto& allocation : allocations) { - if (allocation->memory_space() == MemorySpace::kAlternate) { - MsaBufferInterval buffer_interval; - buffer_interval.start = allocation->start_time(); - buffer_interval.end = allocation->end_time(); - buffer_interval.size = allocation->chunk().size; - buffer_interval.buffer = prefetch_candidate.buffer; - AddToPendingChunks(buffer_interval, chunk_candidate); - for (const HloUse& use : allocation->uses()) { - AddOperandToAlternateMemoryMap(use.instruction, use.operand_number, - use.operand_index); - } - } - allocations_->push_back(std::move(allocation)); - } - - // Add a repack allocation block for the Allocation objects in alternate - // memory. - std::vector colocations; - for (int i = allocations_initial_size; i < allocations_->size(); ++i) { - const auto& allocation = allocations_->at(i); - if (allocation->memory_space() == MemorySpace::kAlternate) { - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - allocation->start_time(), allocation->end_time(), - allocation->chunk().size, allocation->chunk().offset, - allocation.get())); - colocations.push_back(&repack_allocation_blocks_.back()); - } - } - MarkRepackAllocationBlocksColocated(colocations); - ClearPendingChunks(); -} - -void MsaAlgorithm::AllocateReservedScopedAllocations() { - const std::vector& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - std::vector colocations; - for (BreadthFirstMidpointIterator it(0, instruction_sequence.size() - 1); - !it.End(); it.Next()) { - HloInstruction* instruction = instruction_sequence[it.value()]; - int64_t reserved_scoped_memory = - std::min(options_.reserved_scoped_memory_fn( - instruction, /*operands_in_alternate_memory=*/{}, - /*outputs_in_alternate_memory=*/{}), - options_.max_size_in_bytes); - if (reserved_scoped_memory == 0) { - continue; - } - AllocateScopedAllocation(instruction, /*is_post_module=*/false, - reserved_scoped_memory, it.value(), colocations); - } - // If requested, make all scoped allocations to colocate with each other so - // that when we repack, all scoped allocations get the same offsets. Since - // they will all have the same scoped memory addresses, this increases the - // opportunity to deduplicate different ops. However, this may hurt the - // memory packing efficiency. - if (options_.allocate_reserved_scoped_memory_at_same_offset) { - MarkRepackAllocationBlocksColocated(colocations); - } - - colocations.clear(); - // Allocate post-module scoped allocation if requested. It never needs to be - // colocated with other scoped allocations. - if (options_.post_module_scoped_alternate_memory_size_in_bytes > 0) { - AllocateScopedAllocation( - /*instruction=*/module_->entry_computation()->root_instruction(), - /*is_post_module=*/true, - options_.post_module_scoped_alternate_memory_size_in_bytes, - hlo_live_range_.schedule_end_time(), colocations); - } - - ClearPendingChunks(); -} - -void MsaAlgorithm::AllocateScopedAllocation( - HloInstruction* instruction, bool is_post_module, int64_t size, - int64_t time, std::vector& colocations) { - VLOG(1) << "Allocate reserved scoped memory at " << time << " (" - << (is_post_module ? "" : instruction->name()) - << "): " << size; - MsaBufferInterval interval; - interval.buffer = nullptr; - interval.size = size; - interval.start = time; - interval.end = time; - interval.need_allocation = true; - Chunk chunk_candidate = FindChunkCandidate(interval, /*preferred_offset=*/0); - CHECK_EQ(chunk_candidate.offset, 0); - AddToPendingChunks(interval, chunk_candidate); - - if (options_.dump_fn != nullptr || VLOG_IS_ON(3)) { - AppendScopedAllocationBufferInfoDebugString(instruction, time, size, - buffer_info_str_); - } - - allocations_->push_back(std::make_unique( - chunk_candidate, time, instruction, is_post_module)); - - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - time, time, size, - /*initial_offset=*/0, allocations_->back().get())); - repack_allocation_blocks_.back().next_colocated = - &repack_allocation_blocks_.back(); - colocations.push_back(&repack_allocation_blocks_.back()); -} - -int64_t MsaAlgorithm::GetCorrectedUseTime( - const HloInstruction* instruction) const { - const absl::flat_hash_map& schedule = - hlo_live_range_.instruction_schedule(); - if (instruction->opcode() == HloOpcode::kWhile) { - // Given an example while loop and flattened schedule (logical times shown - // on the left): - // - // 0: a = ... - // 1: ... - // cond { - // 2: p = param(0) - // 3: ... - // } - // body { - // 4: p = param(0) - // 5: ... - // 6: ROOT ... - // } - // 7: w = while(a), body=body, cond=cond - // - // When processing "a" (time 0) and its while use (time 7), we update the - // interval to time 0-4. This is so that the remaining interval (5-6) can be - // allocated separately and this buffer doesn't waste alternate memory space - // within the while loop body. - HloComputation* while_body = instruction->while_body(); - // We require while body ROOTs to be the last in the schedule. - CHECK_EQ(schedule.at(while_body->root_instruction()) + 1, - schedule.at(instruction)) - << "While body ROOTs need to be the last in the schedule! " - "Please run RootInstructionSinker."; - // The corrected use is the parameter time. This is so that we can decide on - // alternate memory allocations within the while loop body when we look at - // uses within the while loop body. - return schedule.at(while_body->parameter_instruction(0)); - } - if (instruction->opcode() == HloOpcode::kConditional) { - // The corrected use time is the earliest parameter of the called - // computations. - int64_t use_time = std::numeric_limits::max(); - for (const HloComputation* called_computation : - instruction->called_computations()) { - use_time = std::min( - use_time, schedule.at(called_computation->parameter_instruction(0))); - } - return use_time; - } - // Otherwise, just return the time of the use instruction. - return hlo_live_range_.instruction_schedule().at(instruction); -} - -int64_t MsaAlgorithm::GetCorrectedUseTime(const HloUse& use) const { - return GetCorrectedUseTime(use.instruction); -} - -int64_t MsaAlgorithm::GetExtendedUseTimeIfUseIsView(const HloUse& use) const { - const absl::flat_hash_map& schedule = - hlo_live_range_.instruction_schedule(); - int64_t use_time = schedule.at(use.instruction); - // Only the viewed value itself (operand 0 of the view, on both the read - // and write sides) needs its reservation extended; a view's start index - // operands are consumed at the view's own time. - if (!options_.dus_view_color.has_value() || use.operand_number != 0 || - !use.instruction->shape().has_layout() || - use.instruction->shape().layout().memory_space() != - *options_.dus_view_color) { - return use_time; - } - // The use instruction is a view into the used value; its consumers read the - // value's storage through it at later times. - return std::max(use_time, - ViewExtendedTransitiveUseTime( - use.instruction, *options_.dus_view_color, schedule)); -} - -std::optional -MsaAlgorithm::RequiredMemoryAssignmentAt(const HloValue* buffer, - int64_t time) const { - auto required_assignment_it = required_assignments_.find(buffer); - std::optional required_assignment_at_time; - if (required_assignment_it != required_assignments_.end()) { - for (const RequiredMemoryAssignment& required_assignment : - required_assignment_it->second) { - if (required_assignment.time == time) { - // Sanity check that there is only one required at time. - CHECK(!required_assignment_at_time) - << buffer->ToShortString() << " at time " << time; - required_assignment_at_time = required_assignment; - } - } - } - return required_assignment_at_time; -} - -bool MsaAlgorithm::InefficientSiteRetryCanProgress( - absl::Span sites) const { - // A committed required assignment seen here is always in default memory, - // so an existing requirement means the default memory requirement is - // already in effect: sites colored in alternate memory are never flagged - // as inefficient (they are skipped when collecting inefficient_sites), - // and a committed alternate memory requirement would have CHECK-failed - // the earlier AddRequiredAssignment(kDefault) call that recorded the - // site's requirement in the first place. - absl::flat_hash_set> pending_keys; - for (const auto& [value, assignment] : pending_required_assignments_) { - pending_keys.insert({value, assignment.time}); - } - return absl::c_any_of(sites, [&](const HloPositionOrUse& site) { - const HloValue* value = nullptr; - int64_t time = 0; - if (const auto* position = std::get_if(&site)) { - value = &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - position->instruction, position->index); - time = hlo_live_range_.instruction_schedule().at(position->instruction); - } else { - const HloUse& use = std::get(site); - value = &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - use.instruction->operand(use.operand_number), use.operand_index); - time = GetCorrectedUseTime(use); - } - return !RequiredMemoryAssignmentAt(value, time).has_value() || - pending_keys.contains({value, time}); - }); -} - -std::optional -MsaAlgorithm::RequiredAssignmentForUse(const AllocationValue::Use& use) const { - HloUse hlo_use = use.hlo_use; - const HloValue* value = &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - hlo_use.instruction->operand(hlo_use.operand_number), - hlo_use.operand_index); - int64_t time = GetCorrectedUseTime(hlo_use); - return RequiredMemoryAssignmentAt(value, time); -} - -std::optional -MsaAlgorithm::AliasedRequiredAssignmentForUse( - const AllocationValue::Use& use) const { - std::optional required_assignment; - for (const HloPosition& position : use.aliases) { - const HloValue* value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - position.instruction, position.index); - int64_t time = - hlo_live_range_.instruction_schedule().at(position.instruction); - std::optional required_assignment_for_alias = - RequiredMemoryAssignmentAt(value, time); - if (required_assignment == std::nullopt) { - required_assignment = required_assignment_for_alias; - } else { - CHECK(required_assignment_for_alias == std::nullopt || - required_assignment->memory_space_and_offset_equal( - *required_assignment_for_alias)); - } - } - return required_assignment; -} - -void MsaAlgorithm::AddAliasedRequiredAssignment( - const HloInstruction* instruction, ShapeIndex index, - const Allocation* aliased_allocation, - RequiredMemoryAssignment::Source required_assignment_source) { - AliasedOffset* offset = nullptr; - if (aliased_allocation->memory_space() == MemorySpace::kAlternate) { - offset = GetAliasedOffset(*aliased_allocation); - } - AddRequiredAssignment(instruction, index, aliased_allocation->memory_space(), - required_assignment_source, offset); -} - -void MsaAlgorithm::AddRequiredAssignment( - const HloValue* value, const HloInstruction* instruction, - MemorySpace memory_space, int64_t time, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset, bool add_to_pending) { - RequiredMemoryAssignment required_assignment{memory_space, time, offset, - required_assignment_source}; - CHECK(memory_space != MemorySpace::kDefault || - !IsPositionColoredInAlternateMemoryAtTime(value->defining_position(), - time)) - << "Conflicting required assignment for : " - << value->defining_position().ToString() - << " required assignment: " << required_assignment.ToString(); - - CHECK( - memory_space != MemorySpace::kAlternate || - !IsPositionColoredInDefaultMemoryAtTime(value->defining_position(), time)) - << "Conflicting required assignment for : " - << value->defining_position().ToString() - << " required assignment: " << required_assignment.ToString(); - - // Check for existing required assignment at this time and make sure it is the - // same as this if there is one. - std::optional existing_required_assignment = - RequiredMemoryAssignmentAt(value, time); - - if (existing_required_assignment.has_value()) { - CHECK(required_assignment.memory_space_and_offset_equal( - existing_required_assignment.value())) - << "Failed to add required assignment due to a conflict. Existing " - "required assignment: " - << existing_required_assignment.value().ToString() - << "; new required assignment: " << required_assignment.ToString() - << "; hlo value: " << value->ToShortString(); - VLOG(3) << "Not adding required assignment because there is one already: " - << value->ToShortString() - << " required assignment: " << required_assignment.ToString(); - } else { - VLOG(3) << "Adding required assignment for value : " - << value->ToShortString() - << " required assignment: " << required_assignment.ToString(); - required_assignments_[value].push_back(required_assignment); - if (add_to_pending) { - pending_required_assignments_.push_back({value, required_assignment}); - } - } -} - -void MsaAlgorithm::AddRequiredAssignment( - const HloInstruction* instruction, ShapeIndex index, - MemorySpace memory_space, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset, bool add_to_pending) { - const HloValue* value = - &alias_analysis_.dataflow_analysis().GetUniqueValueAt(instruction, index); - int64_t instruction_time = - hlo_live_range_.instruction_schedule().at(instruction); - AddRequiredAssignment(value, instruction, memory_space, instruction_time, - required_assignment_source, offset, add_to_pending); -} - -void MsaAlgorithm::AddRequiredAssignment( - const HloPosition& position, MemorySpace memory_space, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset, bool add_to_pending) { - AddRequiredAssignment(position.instruction, position.index, memory_space, - required_assignment_source, offset, add_to_pending); -} - -void MsaAlgorithm::AddRequiredAssignment( - const HloUse& use, MemorySpace memory_space, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset, bool add_to_pending) { - const HloValue* value = &alias_analysis_.dataflow_analysis().GetUniqueValueAt( - use.instruction->operand(use.operand_number), use.operand_index); - int64_t instruction_time = GetCorrectedUseTime(use); - AddRequiredAssignment(value, use.instruction, memory_space, instruction_time, - required_assignment_source, offset, add_to_pending); -} - -void MsaAlgorithm::AddInputAndOutputRequiredAssignments() { - // Go through the parameters, outputs, and constants and pin them to the - // corresponding memory by adding a required assignment. - const HloModule& module = alias_analysis_.dataflow_analysis().module(); - const auto& instruction_schedule = hlo_live_range_.instruction_schedule(); - HloComputation* entry_computation = module.entry_computation(); - for (HloInstruction* parameter_instruction : - entry_computation->parameter_instructions()) { - int64_t parameter_instruction_time = - instruction_schedule.at(parameter_instruction); - ShapeUtil::ForEachSubshape( - parameter_instruction->shape(), - [&](const Shape& subshape, const ShapeIndex& index) { - MemorySpace memory_space = MemorySpace::kDefault; - if (subshape.has_layout() && subshape.layout().memory_space() == - options_.alternate_memory_space) { - memory_space = MemorySpace::kAlternate; - } - for (const HloBuffer* buffer : - alias_analysis_.ComputeBuffersAt(parameter_instruction, index)) { - for (const HloValue* value : buffer->values()) { - VLOG(3) << "Adding required assignment for parameter value = " - << value->ToShortString() - << " time = " << parameter_instruction_time << " space = " - << (memory_space == MemorySpace::kDefault ? "def" - : "alt"); - AddRequiredAssignment( - value, parameter_instruction, memory_space, - parameter_instruction_time, - RequiredMemoryAssignment::Source::kProgramInput, - /*offset=*/nullptr, - /*add_to_pending=*/false); - } - } - }); - } - HloInstruction* root_instruction = entry_computation->root_instruction(); - int64_t root_instruction_time = instruction_schedule.at(root_instruction); - ShapeUtil::ForEachSubshape( - root_instruction->shape(), - [&](const Shape& subshape, const ShapeIndex& index) { - MemorySpace memory_space = MemorySpace::kDefault; - if (subshape.has_layout() && subshape.layout().memory_space() == - options_.alternate_memory_space) { - memory_space = MemorySpace::kAlternate; - } - for (const HloBuffer* buffer : - alias_analysis_.ComputeBuffersAt(root_instruction, index)) { - for (const HloValue* value : buffer->values()) { - VLOG(3) << "Adding required assignment for output value = " - << value->ToShortString() - << " time = " << root_instruction_time << " space = " - << (memory_space == MemorySpace::kDefault ? "def" : "alt"); - AddRequiredAssignment( - value, root_instruction, memory_space, root_instruction_time, - RequiredMemoryAssignment::Source::kProgramOutput, - /*offset=*/nullptr, /*add_to_pending=*/false); - } - } - }); - - for (const HloComputation* computation : module.MakeNonfusionComputations()) { - for (HloInstruction* instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::kConstant) { - auto constant_instruction_it = instruction_schedule.find(instruction); - if (constant_instruction_it == instruction_schedule.end()) { - continue; - } - int64_t constant_instruction_time = constant_instruction_it->second; - ShapeUtil::ForEachLeafShape( - instruction->shape(), - [&](const Shape& /*sub_shape*/, const ShapeIndex& index) { - for (const HloBuffer* buffer : - alias_analysis_.ComputeBuffersAt(instruction, index)) { - for (const HloValue* value : buffer->values()) { - VLOG(3) << "Adding required assignment for constant value = " - << value->ToShortString() - << " time = " << constant_instruction_time - << " space = def"; - AddRequiredAssignment( - value, instruction, MemorySpace::kDefault, - constant_instruction_time, - RequiredMemoryAssignment::Source::kConstantInstruction, - /*offset=*/nullptr, - /*add_to_pending=*/false); - } - } - }); - } - } - } - - // Go through all of the values and pin them to the default memory if they are - // not allowed on the alternate memory. - for (const HloValue* value : alias_analysis_.dataflow_analysis().values()) { - if (!options_.is_allowed_in_alternate_mem_fn(*value)) { - // We won't find the instruction in the schedule if it's inside a fusion. - // If so, just skip. - auto instruction_time_it = - instruction_schedule.find(value->instruction()); - if (instruction_time_it == instruction_schedule.end()) { - continue; - } - int64_t instruction_time = instruction_time_it->second; - auto& required_assignments = required_assignments_[value]; - // Check if there is an existing matching required assignment (e.g. - // inserted by the logic above) and if so ensure it requires a default - // memory allocation. - auto matching_assignment = absl::c_find_if( - required_assignments, - [&](const RequiredMemoryAssignment& required_assignment) { - return required_assignment.time == instruction_time; - }); - if (matching_assignment != required_assignments.end()) { - CHECK(matching_assignment->memory_space == MemorySpace::kDefault) - << "Mismatch in required assignments at time " << instruction_time - << " value: " << value->ToString(); - } else { - CHECK(!IsPositionColoredInAlternateMemoryAtTime( - value->defining_position(), instruction_time)) - << "Conflicting input/output required assignment for " - "position: " - << value->defining_position().ToString() << " at " - << instruction_time - << " in default memory, because it is colored in alternate memory"; - VLOG(3) << "Adding required assignment: " << value->ToShortString() - << " at " << instruction_time << " at def"; - required_assignments.push_back( - {MemorySpace::kDefault, instruction_time}); - } - } - } -} - -bool MsaAlgorithm::AreIntervalsReservedInAlternateMemory( - absl::Span colocated_intervals) const { - auto is_position_in_alternate_memory = [&](const HloPosition& position) { - const Shape& shape = position.shape(); - return shape.has_layout() && - shape.layout().memory_space() == options_.alternate_memory_space; - }; - - const HloModule& module = alias_analysis_.dataflow_analysis().module(); - const HloComputation* entry_computation = module.entry_computation(); - const HloInstruction* root_instruction = - entry_computation->root_instruction(); - for (const MsaBufferInterval* colocated_interval : colocated_intervals) { - const HloValue* value = colocated_interval->buffer; - if (value->defining_instruction()->opcode() == HloOpcode::kParameter && - value->defining_instruction()->parent() == entry_computation && - is_position_in_alternate_memory(value->defining_position())) { - return true; - } - - for (const HloPosition& position : value->positions()) { - if (position.instruction == root_instruction && - is_position_in_alternate_memory(position)) { - return true; - } - } - } - return false; -} - -const std::vector* -MsaAlgorithm::GetRepeatedInstructionList( - const HloInstruction* instruction) const { - const auto fingerprint_it = fingerprint_map_.find(instruction); - if (fingerprint_it == fingerprint_map_.end()) { - return nullptr; - } - const auto repeated_insts_it = - repeated_inst_map_.find(fingerprint_it->second); - CHECK(repeated_insts_it != repeated_inst_map_.end()); - return &repeated_insts_it->second; -} - -void MsaAlgorithm::UpdateReservedScopedAllocationSize() { - // Check all instructions, if their operands/outputs have been placed in - // alternate memory, update their scoped allocation size. - VLOG(2) << "Update scoped allocation size before repacking."; - const auto& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - absl::flat_hash_map reserved_scoped_memory_map; - for (int i = 0; i < instruction_sequence.size(); ++i) { - const HloInstruction* instruction = instruction_sequence[i]; - absl::flat_hash_set> - empty_operand_in_alternate_memory; - auto find_operands_in_alternate_memory_it = - operands_in_alternate_memory_map_.find(instruction); - const absl::flat_hash_set>& - operands_in_alternate_memory = - find_operands_in_alternate_memory_it == - operands_in_alternate_memory_map_.end() - ? empty_operand_in_alternate_memory - : find_operands_in_alternate_memory_it->second; - absl::flat_hash_set empty_output_in_alternate_memory; - auto find_outputs_in_alternate_memory_it = - outputs_in_alternate_memory_map_.find(instruction); - const absl::flat_hash_set& outputs_in_alternate_memory = - find_outputs_in_alternate_memory_it == - outputs_in_alternate_memory_map_.end() - ? empty_output_in_alternate_memory - : find_outputs_in_alternate_memory_it->second; - reserved_scoped_memory_map[i] = options_.reserved_scoped_memory_fn( - instruction, operands_in_alternate_memory, outputs_in_alternate_memory); - } - // Update scoped allocation sizes. - for (RepackAllocationBlock& allocation_block : repack_allocation_blocks_) { - Allocation* allocation = allocation_block.allocation; - if (allocation->is_scoped_allocation()) { - int64_t time = allocation->start_time(); - peak_memory_usage_[time] += - (reserved_scoped_memory_map[time] - allocation->chunk().size); - allocation_block.size = reserved_scoped_memory_map[time]; - allocation->mutable_chunk()->size = reserved_scoped_memory_map[time]; - } - } -} - -void MsaAlgorithm::ExportAllocationsForRepacking( - std::vector& allocations) { - using SliceDetail = SlicedCopyAllocation::SliceDetail; - - if (options_.reduce_scoped_memory_limit) { - UpdateReservedScopedAllocationSize(); - } - - for (RepackAllocationBlock& allocation_block : repack_allocation_blocks_) { - allocation_block.original_slice_data = std::nullopt; - allocation_block.repacked_slice_data = std::nullopt; - - if (!allocation_block.allocation->is_sliced_copy_allocation()) { - allocations.push_back(&allocation_block); - continue; - } - - SlicedCopyAllocation* allocation = - dynamic_cast(allocation_block.allocation); - std::vector slice_details_sorted_by_offset; - slice_details_sorted_by_offset.reserve( - allocation->slice_details_sorted_by_start_time().size()); - for (const SliceDetail& slice_detail : - allocation->slice_details_sorted_by_start_time()) { - slice_details_sorted_by_offset.push_back(&slice_detail); - } - absl::c_stable_sort(slice_details_sorted_by_offset, - [](const SliceDetail* lhs, const SliceDetail* rhs) { - return lhs->slice_decision.chunk.offset < - rhs->slice_decision.chunk.offset; - }); - - // Since this is a sliced allocation, construct SlicedAllocationData to - // attach to the AllocationBlock. - SlicedAllocationData original_slice_data; - for (const SliceDetail* slice_detail : slice_details_sorted_by_offset) { - CHECK_EQ(slice_detail->copy_start_after_time, - slice_detail->slice_decision.exclusive_start_time); - original_slice_data.slices_sorted_by_offset.push_back(AllocatedSlice{ - slice_detail->slice_decision.chunk.size, - slice_detail->slice_decision.chunk.offset, - /*inclusive_start_time=*/ - ExclusiveToInclusiveStartTime( - slice_detail->slice_decision.exclusive_start_time)}); - } - - allocation_block.original_slice_data = std::move(original_slice_data); - allocations.push_back(&allocation_block); - } -} - -void MsaAlgorithm::ImportRepackedAllocations() { - interval_tree_.Clear(); - for (RepackAllocationBlock& allocation_block : repack_allocation_blocks_) { - if (allocation_block.allocation->is_sliced_copy_allocation()) { - ImportRepackedSlicedAllocation(allocation_block); - continue; - } - ImportRepackedNonSlicedAllocation(allocation_block); - } -} - -void MsaAlgorithm::ImportRepackedNonSlicedAllocation( - RepackAllocationBlock& block) { - Allocation* allocation = block.allocation; - int64_t original_offset = block.initial_offset; - int64_t repacked_offset = block.offset; - - // Update the Allocation, AllocationBlock, and interval_tree_. - allocation->set_offset(repacked_offset); - block.initial_offset = repacked_offset; - block.offset = -1; - interval_tree_.Add( - block.inclusive_start_time, block.end_time, - HeapSimulator::Chunk::FromOffsetSize(repacked_offset, block.size)); - - VLOG(3) << "Repacking move. offset: " << original_offset << " -> " - << repacked_offset << "; size: " << block.size - << "; Allocation: " << allocation->ToString(); -} - -void MsaAlgorithm::ImportRepackedSlicedAllocation( - RepackAllocationBlock& block) { - using SlicedCopyAllocation = memory_space_assignment::SlicedCopyAllocation; - using SliceDetail = SlicedCopyAllocation::SliceDetail; - - CHECK_OK(AreRepackedSlicesValid(block)); - - SlicedCopyAllocation* allocation = - dynamic_cast(block.allocation); - CHECK(block.allocation->is_sliced_copy_allocation()); - int64_t original_offset = block.initial_offset; - int64_t repacked_offset = block.offset; - std::vector original_slice_offsets = - allocation->SliceOffsetsSortedByStartTime(); - - // Update the Allocation, AllocationBlock, and interval_tree_. - allocation->set_offset(repacked_offset); - if (block.repacked_slice_data.has_value()) { - allocation->ImportRepackedSliceData(*block.repacked_slice_data); - } else { - allocation->AddDiffToAllSliceOffsets(repacked_offset - original_offset); - } - block.initial_offset = repacked_offset; - block.offset = -1; - // Note, in a non-repacking setting, we would have reworked the chunks as - // described in - // MsaAlgorithm::PrefetchContext::SlicedSolution::slices_for_pending_chunks. - // Doing so was for the benefit of MsaAlgorithm::pending_chunks_. However, - // pending_chunks_ are cleared before repacking, when UncommitPendingWork() - // is called. Thus, we don't need to worry about modifying the chunks here. - for (const SliceDetail& slice_detail : - allocation->slice_details_sorted_by_start_time()) { - interval_tree_.Add( - /*start=*/ - ExclusiveToInclusiveStartTime(slice_detail.copy_start_after_time), - block.end_time, slice_detail.slice_decision.chunk); - } - - VLOG(3) << "Repacking move. offset: " << original_offset << " -> " - << repacked_offset << "; size: " << block.size << "; " << - [&]() { - std::vector new_slice_offsets = - allocation->SliceOffsetsSortedByStartTime(); - CHECK_EQ(original_slice_offsets.size(), new_slice_offsets.size()); - std::vector offset_moves; - offset_moves.reserve(original_slice_offsets.size()); - for (int i = 0; i < original_slice_offsets.size(); ++i) { - offset_moves.push_back(absl::StrCat(original_slice_offsets[i], " -> ", - new_slice_offsets[i])); - } - return absl::StrCat("slice_offsets: [", - absl::StrJoin(offset_moves, ", "), "]"); - }() - << "; Allocation: " << allocation->ToString(); -} - -absl::Status MsaAlgorithm::AreRepackedSlicesValid( - const RepackAllocationBlock& block) { - if (!block.repacked_slice_data.has_value()) { - return absl::OkStatus(); - } - if (!block.original_slice_data.has_value()) { - return InvalidArgumentStrCat( - "Repacked sliced allocation has repacked slice data but not original " - "slice data."); - } - int64_t num_slices = - block.original_slice_data->slices_sorted_by_offset.size(); - if (num_slices != block.repacked_slice_data->slices_sorted_by_offset.size()) { - return InvalidArgumentStrCat( - "Repacked sliced allocation has ", num_slices, - " slices but repacking has data for ", - block.repacked_slice_data->slices_sorted_by_offset.size(), " slices."); - } - - // Ensure that the slice size to start time mapping has not changed. If it - // changes, its invalidates MSA's internal state, e.g., the peak_memory_usage_ - // data structure. - std::vector> original_size_to_time_mapping; - original_size_to_time_mapping.reserve(num_slices); - for (const AllocatedSlice& slice : - block.original_slice_data->slices_sorted_by_offset) { - original_size_to_time_mapping.push_back( - std::make_pair(slice.size, slice.inclusive_start_time)); - }; - absl::c_sort(original_size_to_time_mapping); - std::vector> repacked_size_to_time_mapping; - repacked_size_to_time_mapping.reserve(num_slices); - for (const AllocatedSlice& slice : - block.repacked_slice_data->slices_sorted_by_offset) { - repacked_size_to_time_mapping.push_back( - std::make_pair(slice.size, slice.inclusive_start_time)); - }; - absl::c_sort(repacked_size_to_time_mapping); - if (original_size_to_time_mapping != repacked_size_to_time_mapping) { - return InvalidArgumentStrCat( - "Repacked slices do not preserve the initial slice size-start time " - "mappings."); - } - - return absl::OkStatus(); -} - -bool MsaAlgorithm::UncommitChunkAndUpdatePeakMemory( - const MsaBufferInterval& interval, const Chunk& chunk) { - VLOG(3) << "Uncommitting: [" << interval.start << ", " << interval.end - << "] off = " << chunk.offset << " size = " << chunk.size; - for (int i = interval.start; i <= interval.end; ++i) { - peak_memory_usage_[i] -= chunk.size; - CHECK_GE(peak_memory_usage_[i], 0) - << "Peak memory usage at " << i << " is below zero after uncommitting. " - << interval.start << "-" << interval.end << " : [" << chunk.offset - << ", " << chunk.size << "]"; - } - return interval_tree_.Remove(interval.start, interval.end, chunk); -} - -void MsaAlgorithm::CommitChunkAndUpdatePeakMemory( - const MsaBufferInterval& buffer_interval, const Chunk& chunk) { - for (int i = buffer_interval.start; i <= buffer_interval.end; ++i) { - peak_memory_usage_[i] += chunk.size; - CHECK_LE(peak_memory_usage_[i], options_.max_size_in_bytes) - << "Peak memory usage at " << i - << " exceeds the max size of alternate memory. " - << buffer_interval.start << "-" << buffer_interval.end << " : " - << chunk.ToString() << " buffer " - << (buffer_interval.buffer ? buffer_interval.buffer->ToString() - : "null"); - } - CommitChunkAndInterval(buffer_interval, chunk); -} - -void MsaAlgorithm::UncommitPendingWork( - absl::Span allocation_values) { - // Clear the allocation sequence of the allocation values so that in case we - // retry allocation after uncommitting. - for (AllocationValue& allocation_value : allocation_values) { - allocation_value.mutable_allocation_sequence()->clear(); - } - for (const auto& interval_and_chunk : pending_chunks_) { - UncommitChunkAndUpdatePeakMemory(interval_and_chunk.first, - interval_and_chunk.second); - } - for (const AsynchronousCopy& async_copy : pending_async_copies_) { - if (async_copy.destination == MemorySpace::kAlternate) { - prefetch_interval_tree_.Remove( - /*start=*/ - ExclusiveToInclusiveStartTime(async_copy.exclusive_start_time), - async_copy.end_time, kDummyChunk); - prefetch_async_copy_resource_.RemoveCopy(async_copy); - if (options_.enforce_prefetch_fifo_order) { - async_copy_ordering_.RemoveCopy(async_copy); - } - } else { - eviction_interval_tree_.Remove( - /*start=*/ - ExclusiveToInclusiveStartTime(async_copy.exclusive_start_time), - async_copy.end_time, kDummyChunk); - eviction_async_copy_resource_.RemoveCopy(async_copy); - } - } - for (const auto& value_and_required_assignment : - pending_required_assignments_) { - auto& required_assignment_vector = - required_assignments_[value_and_required_assignment.first]; - const RequiredMemoryAssignment& required_assignment = - value_and_required_assignment.second; - VLOG(3) << "Removing required assignment: " - << (required_assignment.memory_space == MemorySpace::kDefault - ? "def" - : "alt") - << " time = " << required_assignment.time << " off = " - << (required_assignment.offset ? required_assignment.offset->offset - : -1); - for (auto it = required_assignment_vector.begin(); - it != required_assignment_vector.end(); ++it) { - if (*it == value_and_required_assignment.second) { - required_assignment_vector.erase(it); - break; - } - } - } - // As part of the uncommit, we need to re-reserve the previously freed chunks, - // corresponding to the deallocated reserved allocations that were pending. - for (const auto& allocation : pending_deallocated_reserved_allocations_) { - MsaBufferInterval interval = - MsaBufferInterval{/*buffer=*/nullptr, - /*size=*/allocation->chunk().size, - /*start=*/allocation->start_time(), - /*end=*/allocation->end_time(), - /*colocations=*/{}, - /*need_allocation=*/true}; - Chunk chunk_candidate = FindChunkCandidate( - interval, /*preferred_offset=*/allocation->chunk().offset); - CHECK_EQ(chunk_candidate.offset, allocation->chunk().offset); - AddToPendingChunks(interval, chunk_candidate); - allocation->mark_chunk_reserved_in_interval_tree(); - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - allocation->start_time(), allocation->end_time(), chunk_candidate.size, - chunk_candidate.offset, allocation)); - repack_allocation_blocks_.back().next_colocated = - &(repack_allocation_blocks_.back()); - } - ClearPendingChunks(); -} - -void MsaAlgorithm::FinalizeAllocations( - absl::Span allocation_values) { - VLOG(3) << "Finalized allocations:"; - for (const HloInstruction* copy_inst : sorted_async_conversion_candidates_) { - successful_async_conversion_set_.insert(copy_inst); - } - not_finalized_async_conversions_.clear(); - std::vector>> - colocation_vector; - absl::flat_hash_map offset_to_index; - for (int allocation_value_idx = 0; - allocation_value_idx < allocation_values.size(); - ++allocation_value_idx) { - auto& allocation_value = allocation_values.at(allocation_value_idx); - VLOG(3) << " " << allocation_value_idx + 1 << "/" - << allocation_values.size() << ") " - << allocation_value.ToShortString(); - for (auto& allocation : *allocation_value.mutable_allocation_sequence()) { - VLOG(3) << " " << allocation->ToString(); - 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()) { - AddOperandToAlternateMemoryMap(use.instruction, use.operand_number, - use.operand_index); - } - if (!allocation->is_copy_like_allocation()) { - outputs_in_alternate_memory_map_[allocation->defining_position() - .instruction] - .insert(allocation->defining_position().index); - } - } - allocations_->push_back(std::move(allocation)); - Allocation* inserted_allocation = allocations_->back().get(); - if (inserted_allocation->memory_space() == MemorySpace::kAlternate) { - auto* aliased_offset = GetAliasedOffset(*inserted_allocation); - auto [it, inserted] = - offset_to_index.emplace(aliased_offset, colocation_vector.size()); - if (inserted) { - colocation_vector.emplace_back(aliased_offset, - std::vector{inserted_allocation}); - } else { - size_t index = it->second; - colocation_vector[index].second.push_back(inserted_allocation); - } - // Register the allocation start time - 1 and end time + 1 as the edge - // time indices for the allocation to indicate the time indices that the - // alternate memory will be freer than the next index. Evict uses this - // to bypass some of the redundant calls to FindChunkCandidate. - edge_time_indices_.insert(inserted_allocation->start_time() - 1); - edge_time_indices_.insert(inserted_allocation->end_time() + 1); - } - } - } - // The allocations that have the same AliasedOffset need to be colocated. - // Export these to repack_allocation_blocks_ so that we can repack them to - // reduce fragmentation. - for (auto& colocation : colocation_vector) { - std::vector colocations; - for (Allocation* colocated_allocation : colocation.second) { - repack_allocation_blocks_.push_back(MakeRepackAllocationBlock( - colocated_allocation->start_time(), colocated_allocation->end_time(), - colocated_allocation->chunk().size, - colocated_allocation->chunk().offset, colocated_allocation)); - colocations.push_back(&repack_allocation_blocks_.back()); - } - MarkRepackAllocationBlocksColocated(colocations); - } - ClearPendingChunks(); -} - -void MsaAlgorithm::ClearPendingChunks() { - pending_chunks_.clear(); - pending_async_copies_.clear(); - pending_required_assignments_.clear(); - aliased_offset_map_.clear(); - aliased_offsets_.clear(); - // Ensure auxiliary maps referencing raw AliasedOffset* pointers owned by - // aliased_offsets_ are cleared during uncommits to prevent dangling pointer - // segmentation faults upon subsequent allocation retries. - pipelined_while_preferred_offset_for_computation_.clear(); - pipelined_while_buffer_id_to_aliased_offset_.clear(); - pending_deallocated_reserved_allocations_.clear(); -} - -bool MsaAlgorithm::IsInstructionPendingReplacements( - const HloInstruction* instruction) const { - return std::find(sorted_async_conversion_candidates_.begin(), - sorted_async_conversion_candidates_.end(), - instruction) != sorted_async_conversion_candidates_.end(); -} - -void MsaAlgorithm::AddToPendingChunks(const MsaBufferInterval& buffer_interval, - const Chunk& chunk) { - VLOG(3) << "Committing chunk: " << buffer_interval.start << "-" - << buffer_interval.end << " : " << chunk.ToString(); - CommitChunkAndUpdatePeakMemory(buffer_interval, chunk); - pending_chunks_.emplace_back(buffer_interval, chunk); -} - -std::optional MsaAlgorithm::FindEarliestExclusiveTimeToSatisfyPeakMemory( - int exclusive_start_time, int end_time, int64_t size) const { - std::optional earliest_time_exclusive = std::nullopt; - for (int time_inclusive = ExclusiveToInclusiveEndTime(end_time); - time_inclusive > exclusive_start_time; --time_inclusive) { - if (peak_memory_usage_[time_inclusive] + size <= - options_.max_size_in_bytes) { - earliest_time_exclusive = InclusiveToExclusiveStartTime(time_inclusive); - } else { - break; - } - } - - return earliest_time_exclusive; -} - -std::string MsaAlgorithm::SingleFailureResultToString( - const AllocationResult& result) { - switch (result) { - case AllocationResult::kSuccess: - return "Success"; - case AllocationResult::kFailOutOfMemory: - return "FailOutOfMemory"; - case AllocationResult::kFailPrevAllocationNotInAlternateMem: - return "FailPrevAllocationNotInAlternateMem"; - case AllocationResult::kFailLiveRangeTooLong: - return "FailLiveRangeTooLong"; - case AllocationResult::kFailLiveRangeTooShort: - return "FailLiveRangeTooShort"; - case AllocationResult::kFailOutOfAsyncCopies: - return "FailOutOfAsyncCopies"; - case AllocationResult::kFailViolatesAsyncCopyResource: - return "FailViolatesAsyncCopyResource"; - case AllocationResult::kFailRequiresUncommit: - return "FailRequiresUncommit"; - case AllocationResult::kAllSlicesHaveTheSameStartTime: - return "AllSlicesHaveTheSameStartTime"; - case AllocationResult::kFailConflictingPreferredOffsets: - return "FailConflictingPreferredOffsets"; - case AllocationResult::kFailSyncDataMoveReplacement: - return "FailSyncDataMoveReplacement"; - default: - return "UnknownResult"; - } -} - -std::string MsaAlgorithm::ResultToString(const AllocationResult& result) { - if (result == AllocationResult::kSuccess) { - return "Success"; - } - std::string result_str = ""; - for (int failure_order = 0; failure_order < 16; ++failure_order) { - AllocationResult failure_value = - static_cast(1 << failure_order); - if (result_is(result, static_cast(failure_value))) { - result_str += (SingleFailureResultToString(failure_value) + " | "); - } - } - result_str = result_str.substr(0, result_str.size() - 3); - return result_str; -} - -void MsaAlgorithm::CheckAndUpdateForDualLiveAllocationValues( - const std::optional& - required_memory_assignment_at_start, - AllocationRequest& request) { - if (!request.allocation_value->requires_contiguous_allocation()) { - return; - } - if (!required_memory_assignment_at_start.has_value()) { - return; - } - if (required_memory_assignment_at_start->memory_space != - MemorySpace::kAlternate) { - return; - } - // Go through previous allocations, for the same HloValue, and check if they - // have already allocated alternate memory at the beginning of the current - // AllocationValue, such that we are required to use the same heap offset. - std::vector overlapping_allocations; - Chunk required_chunk = Chunk::FromOffsetSize( - required_memory_assignment_at_start->offset->offset, request.size); - for (const AllocationValue& processed_allocation_value : - request.processed_allocation_values) { - for (const std::unique_ptr& allocation : - *processed_allocation_value.allocation_sequence()) { - if (allocation->is_in_alternate_mem() && - allocation->start_time() <= request.inclusive_start_time && - request.inclusive_start_time <= allocation->end_time() && - allocation->chunk() == required_chunk) { - overlapping_allocations.push_back(allocation.get()); - } - } - } - absl::c_sort(overlapping_allocations, - [](const Allocation* a, const Allocation* b) { - return a->start_time() < b->start_time(); - }); - int64_t chunk_start_time = request.inclusive_start_time; - for (const Allocation* allocation : overlapping_allocations) { - chunk_start_time = std::max(chunk_start_time, allocation->end_time() + 1); - } - - // Note, we don't have to set request.preferred_offset, or do anything special - // to handle aliasing. This is done for us. Specifically, before calling - // CheckAndUpdateForDualLiveAllocationValues(), AllocateSegment() inserts a - // PinnedAllocation with no associated heap chunk, at the beginning of - // request.allocation_value. It aliases that PinnedAllocation with any - // overlapping allocations calculated above. In - // AllocateInAlternateMemoryNoCopy(), we will find that PinnedAllocation and - // realize we need to use the same alternate memory offset. - request.no_copy_chunk_inclusive_start_time = chunk_start_time; - VLOG(3) << "Setting the no-copy chunk (inc) start time to " - << chunk_start_time; -} - -void MsaAlgorithm::ReleaseReservedAllocationForAlternateMemoryColorings( - ReservedAllocation* reserved_allocation) { - if (!reserved_allocation->is_chunk_reserved_in_interval_tree()) { - return; - } - MsaBufferInterval buffer_interval = MsaBufferInterval{ - /*buffer=*/nullptr, - /*size=*/reserved_allocation->chunk().size, - /*start=*/reserved_allocation->start_time(), - /*end=*/reserved_allocation->end_time(), - /*colocations=*/{}, - /*need_allocation=*/true, - }; - // Release the reserved chunk from the interval tree. - CHECK(UncommitChunkAndUpdatePeakMemory(buffer_interval, - reserved_allocation->chunk())); - reserved_allocation->mark_chunk_freed_in_interval_tree(); - pending_deallocated_reserved_allocations_.push_back(reserved_allocation); - size_t original_size = repack_allocation_blocks_.size(); - repack_allocation_blocks_.remove_if( - [reserved_allocation]( - const RepackAllocationBlock& repack_allocation_block) { - return repack_allocation_block.allocation == reserved_allocation; - }); - CHECK_EQ(original_size - repack_allocation_blocks_.size(), 1); -} - -void MsaAlgorithm::FreeAlternateMemoryColoringReservedAllocations( - AllocationRequest& request) { - if (!request.require_start_colored_in_alternate_memory && - !request.require_end_colored_in_alternate_memory) { - return; - } - const HloPosition& defining_position = - request.allocation_value_to_update->position(); - const HloBuffer& buffer = alias_analysis_.GetUniqueBufferAt( - defining_position.instruction, defining_position.index); - auto reserved_allocations_it = - reserved_allocations_for_alt_mem_colorings_.find(&buffer); - CHECK(reserved_allocations_it != - reserved_allocations_for_alt_mem_colorings_.end()); - - int64_t inclusive_start_time = request.inclusive_start_time; - int64_t use_time = request.end_time; - for (std::unique_ptr& reserved_allocation_ptr : - reserved_allocations_it->second) { - if (request.require_start_colored_in_alternate_memory && - reserved_allocation_ptr->start_time() <= inclusive_start_time && - inclusive_start_time <= reserved_allocation_ptr->end_time()) { - ReleaseReservedAllocationForAlternateMemoryColorings( - reserved_allocation_ptr.get()); - } - if (request.require_end_colored_in_alternate_memory && - reserved_allocation_ptr->start_time() <= use_time && - use_time <= reserved_allocation_ptr->end_time()) { - ReleaseReservedAllocationForAlternateMemoryColorings( - reserved_allocation_ptr.get()); - } - } -} - -void MsaAlgorithm::UpdateRequestWithAlternateMemoryColoringRequirements( - AllocationRequest& request) { - if (!request.allocation_value) { - return; - } - const HloPosition& defining_position = - request.allocation_value_to_update->position(); - - int64_t inclusive_start_time = request.inclusive_start_time; - int64_t use_time = request.end_time; - - const HloBuffer& buffer = alias_analysis_.GetUniqueBufferAt( - defining_position.instruction, defining_position.index); - auto reserved_allocations_it = - reserved_allocations_for_alt_mem_colorings_.find(&buffer); - if (reserved_allocations_it != - reserved_allocations_for_alt_mem_colorings_.end()) { - for (std::unique_ptr& reserved_allocation_ptr : - reserved_allocations_it->second) { - if (reserved_allocation_ptr->start_time() <= inclusive_start_time && - inclusive_start_time <= reserved_allocation_ptr->end_time()) { - request.require_start_colored_in_alternate_memory = true; - } - if (reserved_allocation_ptr->start_time() <= use_time && - use_time <= reserved_allocation_ptr->end_time()) { - request.require_end_colored_in_alternate_memory = true; - } - } - } -} - -void MsaAlgorithm::UpdateRequestWithDefaultMemoryColoringRequirements( - AllocationRequest& request) { - if (!request.allocation_value) { - return; - } - const HloPosition& defining_position = - request.allocation_value_to_update->position(); - - int64_t inclusive_start_time = request.inclusive_start_time; - int64_t use_time = request.end_time; - - const HloBuffer& buffer = alias_analysis_.GetUniqueBufferAt( - defining_position.instruction, defining_position.index); - auto default_memory_colorings_it = - default_memory_coloring_requirements_.find(&buffer); - if (default_memory_colorings_it != - default_memory_coloring_requirements_.end()) { - for (const TimeInterval& time_interval : - default_memory_colorings_it->second) { - if (time_interval.inclusive_start_time <= inclusive_start_time && - inclusive_start_time <= time_interval.inclusive_end_time) { - request.require_start_colored_in_default_memory = true; - } - if (time_interval.inclusive_start_time <= use_time && - use_time <= time_interval.inclusive_end_time) { - request.require_end_colored_in_default_memory = true; - } - } - } -} - -AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { - auto allocation_sequence = - request.allocation_value->mutable_allocation_sequence(); - // inclusive_start_time == end_time is a special case where the value is - // consumed multiple times by the same instruction. We can just find the - // previous allocation and use that allocation. - if (request.inclusive_start_time == request.end_time) { - Allocation* allocation = GetLiveAllocationAt( - *request.allocation_value_to_update->mutable_allocation_sequence(), - request.end_time); - CHECK_NE(allocation, nullptr); - allocation->AddUse(request.use->hlo_use); - return AllocationResult::kSuccess; - } - - const HloPosition& defining_position = request.allocation_value->position(); - VLOG(2) << "Finding allocation for " - << request.allocation_value->ToShortString() << " [" - << request.inclusive_start_time << ", " << request.end_time - << "] latest prefetch = " << request.latest_prefetch_time - << " last use = " << request.allocation_value->uses().back().time - << " use = " << request.use->hlo_use.ToString() - << ". Size = " << request.size - << ", def pos = " << defining_position.ToString(); - if (request.require_no_copy_alternate_mem_allocation) { - VLOG(2) << "Requiring alternate memory allocation."; - } - CHECK_LE(request.inclusive_start_time, request.end_time); - if (VLOG_IS_ON(3) && options_.cost_analysis) { - const HloPosition& defining_position = request.allocation_value->position(); - const HloUse& use = request.use->hlo_use; - VLOG(3) << "Definition benefit = " - << options_.cost_analysis->GetAlternateMemoryBenefit( - request.allocation_value->position()) - << " use benefit = " - << options_.cost_analysis->GetAlternateMemoryBenefit( - request.use->hlo_use); - VLOG(3) << "Definition bytes accessed = " - << options_.cost_analysis->OutputBytesAccessed( - *defining_position.instruction, defining_position.index) - << ", use bytes accessed = " - << options_.cost_analysis->OperandBytesAccessed( - *use.instruction, use.operand_number, use.operand_index); - } - - // There could be a requirement to pin this buffer to default memory either - // because it is a parameter or an output. If the buffer is a parameter, then - // we're allowed to prefetch. If the use expects the output to be in default - // memory, we cannot prefetch it because if we did, it would be in alternate - // memory instead. - std::optional required_assignment_at_start = - RequiredMemoryAssignmentAt(request.allocation_value->value(), - request.inclusive_start_time); - std::optional required_memory_space_at_start; - if (required_assignment_at_start.has_value()) { - if (request.require_no_copy_alternate_mem_allocation && - required_assignment_at_start->memory_space == MemorySpace::kDefault) { - required_assignment_at_start = std::nullopt; - } else { - required_memory_space_at_start = - required_assignment_at_start->memory_space; - } - } - // Find required assignment both for the use and its aliases. If they are both - // non-nullopt, then make sure they require the same assignment. - std::optional required_assignment_at_end = - RequiredMemoryAssignmentAt(request.allocation_value_to_update->value(), - request.end_time); - std::optional aliased_required_assignment_at_end = - AliasedRequiredAssignmentForUse(*request.use); - if (required_assignment_at_end != aliased_required_assignment_at_end) { - if (required_assignment_at_end == std::nullopt) { - required_assignment_at_end = aliased_required_assignment_at_end; - } else { - CHECK(aliased_required_assignment_at_end == std::nullopt || - aliased_required_assignment_at_end->memory_space_and_offset_equal( - *required_assignment_at_end)) - << "Conflicting aliased required assignment at end." - " required_assignment_at_end: " - << required_assignment_at_end->ToString() - << " aliased_required_assignment_at_end: " - << aliased_required_assignment_at_end->ToString() - << " for alised use: " << request.use->hlo_use.ToString(); - } - } - std::optional required_memory_space_at_end; - if (required_assignment_at_end.has_value()) { - required_memory_space_at_end = required_assignment_at_end->memory_space; - } - - CHECK(!required_memory_space_at_end.has_value() || - required_memory_space_at_end != MemorySpace::kAlternate || - !request.require_end_colored_in_default_memory) - << "End colored in default memory for an allocation requiring " - "alternate memory space at end. Allocation: " - << request.allocation_value->ToShortString() - << " Use: " << request.use->hlo_use.ToString() - << " start time: " << request.inclusive_start_time - << " end time: " << request.end_time; - CHECK(!required_memory_space_at_end.has_value() || - required_memory_space_at_end != MemorySpace::kDefault || - !request.require_end_colored_in_alternate_memory) - << "End colored in alternate memory for an allocation requiring " - "default memory space at end. Allocation: " - << request.allocation_value->ToShortString() - << " Use: " << request.use->hlo_use.ToString() - << " start time: " << request.inclusive_start_time - << " end time: " << request.end_time; - CHECK(!required_memory_space_at_start.has_value() || - required_memory_space_at_start != MemorySpace::kAlternate || - !request.require_start_colored_in_default_memory) - << "Start colored in default memory for an allocation requiring " - "alternate memory space at start. Allocation:" - << request.allocation_value->ToShortString() - << " Use: " << request.use->hlo_use.ToString() - << " start time: " << request.inclusive_start_time - << " end time: " << request.end_time; - CHECK(!required_memory_space_at_start.has_value() || - required_memory_space_at_start != MemorySpace::kDefault || - !request.require_start_colored_in_alternate_memory) - << "Start colored in alternate memory for an allocation requiring " - "default memory space at start. Allocation:" - << request.allocation_value->ToShortString() - << " Use: " << request.use->hlo_use.ToString() - << " start time: " << request.inclusive_start_time - << " end time: " << request.end_time; - - if (required_assignment_at_start.has_value()) { - bool needs_required_allocation = true; - if (!allocation_sequence->empty()) { - auto prev_allocation_it = std::find_if( - allocation_sequence->rbegin(), allocation_sequence->rend(), - [&](const auto& allocation) { - return allocation->memory_space() == required_memory_space_at_start; - }); - if (prev_allocation_it != allocation_sequence->rend()) { - (*prev_allocation_it)->Extend(request.inclusive_start_time); - needs_required_allocation = false; - } - } - if (needs_required_allocation) { - std::optional aliased_chunk = std::nullopt; - if (required_assignment_at_start->memory_space == - MemorySpace::kAlternate) { - aliased_chunk = Chunk::FromOffsetSize( - required_assignment_at_start->offset->offset, request.size); - } - allocation_sequence->push_back(std::make_unique( - defining_position, required_assignment_at_start->memory_space, - aliased_chunk, request.inclusive_start_time, - request.inclusive_start_time)); - if (required_assignment_at_start->memory_space == - MemorySpace::kAlternate) { - MaybeCreateOrAddToAliasedOffset(*allocation_sequence->back(), - required_assignment_at_start->offset); - } - } - } - - VLOG(3) - << "Required memory assignment at start: " - << OptionalRequiredMemoryAssignmentToString(required_assignment_at_start) - << "; required memory assignment at end: " - << OptionalRequiredMemoryAssignmentToString(required_assignment_at_end); - - FreeAlternateMemoryColoringReservedAllocations(request); - AllocationResult allocation_result = AllocationResult::kSuccess; - // First try keeping the allocation entirely in the alternate memory. - if (!request.require_start_colored_in_default_memory && - !request.require_end_colored_in_default_memory && - (request.require_no_copy_alternate_mem_allocation || - (required_memory_space_at_start != MemorySpace::kDefault && - required_memory_space_at_end != MemorySpace::kDefault)) && - request.allow_no_copy_alternate_mem_allocation && - !request.require_copy_allocation) { - CheckAndUpdateForDualLiveAllocationValues(required_assignment_at_start, - request); - allocation_result = AllocateInAlternateMemoryNoCopy(request); - if (allocation_result == AllocationResult::kSuccess) { - return AllocationResult::kSuccess; - } - // If we required alternate memory allocation, return on failure. - if (request.require_no_copy_alternate_mem_allocation) { - return allocation_result; - } - } - - if (request.require_no_copy_alternate_mem_allocation) { - return allocation_result != AllocationResult::kSuccess - ? allocation_result - : AllocationResult::kFailOutOfMemory; - } - - if (request.require_start_colored_in_alternate_memory) { - // Since no-copy-allocation failed, continuous allocation is not possible in - // the alternate memory. - if (request.allocation_value->requires_contiguous_allocation()) { - return allocation_result; - } - if (request.allocation_value->allocation_sequence()->empty()) { - allocation_result = ForceAlternateMemoryAllocationForMinTime(request); - // We only allocate in alternate memory for one logical time unit by - // adding an immediate eviction later. Allocation for one logical time - // unit should always succeed since we released a reserved chunk that was - // a fallback for this purpose. - CHECK(allocation_result == AllocationResult::kSuccess); - } else if (!IsAsyncConversionCandidate( - request.allocation_value_to_update->position() - .instruction)) { - // If the start of an allocation is in alternate memory and the allocation - // sequence is not empty: - // - It is either an async conversion candidate, which means the coloring - // requirement will be fulfilled with a copy allocation - // - Or we already have an allocation in the alternate memory. - Allocation* last_allocation = - request.allocation_value->allocation_sequence()->back().get(); - CHECK(last_allocation->memory_space() == MemorySpace::kAlternate && - last_allocation->end_time() >= request.inclusive_start_time); - } - } - - auto prev_allocation_it = allocation_sequence->rbegin(); - // Find a previous allocation that is in the default memory space (not - // necessarily the very last allocation). - auto prev_allocation_in_default_mem_it = - std::find_if(allocation_sequence->rbegin(), allocation_sequence->rend(), - [&](const auto& allocation) { - return allocation->memory_space() == MemorySpace::kDefault; - }); - - if (!request.allocation_value->requires_contiguous_allocation()) { - if (prev_allocation_in_default_mem_it == allocation_sequence->rend() && - prev_allocation_it != allocation_sequence->rend() && - (*prev_allocation_it)->memory_space() == MemorySpace::kAlternate && - (*prev_allocation_it)->defining_position() == defining_position) { - // If there was an allocation for this HloValue that was in the alternate - // memory space, we also need to perform an eviction. - // If start or end is required in alternate memory we need to force evict - // since the last allocation in alternate memory could not be extended. - AllocationResult eviction_result = Evict( - request, - /*force_evict=*/request.require_start_colored_in_alternate_memory || - request.require_end_colored_in_alternate_memory); - if (eviction_result != AllocationResult::kSuccess) { - // A non-success eviction requires us to uncommit previous allocations. - return result_mark(AllocationResult::kFailRequiresUncommit, - eviction_result); - } - prev_allocation_in_default_mem_it = allocation_sequence->rbegin(); - } else if (prev_allocation_in_default_mem_it == - allocation_sequence->rend()) { - allocation_sequence->push_back(std::make_unique( - defining_position, MemorySpace::kDefault, - /*chunk=*/std::nullopt, request.inclusive_start_time, - request.end_time)); - prev_allocation_in_default_mem_it = allocation_sequence->rbegin(); - } - } else if (prev_allocation_it == allocation_sequence->rend() && - (request.require_start_colored_in_default_memory || - request.require_end_colored_in_default_memory)) { - // There are no previous allocations, we require contiguous allocation and - // either the start or end needs to be colored in the default memory. - // We can satisfy this requirement by pinning the allocation in the default - // memory space for this time range. - allocation_sequence->push_back(std::make_unique( - defining_position, MemorySpace::kDefault, - /*chunk=*/std::nullopt, request.inclusive_start_time, - request.end_time)); - prev_allocation_in_default_mem_it = allocation_sequence->rbegin(); - } else if (prev_allocation_in_default_mem_it == allocation_sequence->rend()) { - VLOG(3) << "Allocation requires contiguous allocation, but it wasn't " - "possible to find one in alternate memory or default memory."; - return result_mark(AllocationResult::kFailRequiresUncommit, - allocation_result); - } - - CHECK(prev_allocation_in_default_mem_it != allocation_sequence->rend()); - CHECK((*prev_allocation_in_default_mem_it)->memory_space() == - MemorySpace::kDefault); - - // If the allocation value requires a contiguous allocation but has a memory - // space mismatch between the start and end required assignments, then we need - // to uncommit. - if (request.allocation_value->requires_contiguous_allocation() && - required_memory_space_at_start.has_value() && - required_memory_space_at_end.has_value() && - required_memory_space_at_start != required_memory_space_at_end) { - VLOG(3) << "Allocation requires contiguous allocation but has memory space " - "mismatch."; - return result_mark(AllocationResult::kFailRequiresUncommit, - allocation_result); - } - - // Finally, try to prefetch the buffer into alternate memory. - if ((request.allow_prefetch || - request.require_end_colored_in_alternate_memory) && - !request.allocation_value->requires_contiguous_allocation() && - !request.only_extend_existing_allocation && - required_memory_space_at_end != MemorySpace::kDefault && - !request.require_end_colored_in_default_memory) { - if (request.require_copy_allocation && !request.required_copy_for_slice) { - auto it = std::find_if( - allocation_sequence->begin(), allocation_sequence->end(), - [&](const std::unique_ptr& - allocation_ptr) { - auto copy_allocation = - dynamic_cast(allocation_ptr.get()); - return copy_allocation && - copy_allocation->copy_done_schedule_before() <= - request.required_copy_allocation_latest_time; - }); - if (it == allocation_sequence->end()) { - int64_t latest_prefetch_time = - std::min(request.latest_prefetch_time, - request.required_copy_allocation_latest_time); - VLOG(3) << "Updating the latest prefetch time from " - << request.latest_prefetch_time << " to " - << latest_prefetch_time - << ", because this use requires a copy allocation before " - << request.required_copy_allocation_latest_time; - request.latest_prefetch_time = latest_prefetch_time; - } - } - AllocationResult prefetch_result = - Prefetch(request, **prev_allocation_in_default_mem_it, - /*force_prefetch=*/ - request.require_end_colored_in_alternate_memory); - if (prefetch_result == AllocationResult::kSuccess) { - if (request.preferred_prefetch_time) { - // Warn if the prefetch time picked doesn't match the preferred prefetch - // time. - CHECK(!request.allocation_value_to_update->allocation_sequence() - ->empty()); - const Allocation* allocation = - request.allocation_value_to_update->allocation_sequence() - ->back() - .get(); - int64_t prefetch_time = 0; - if (allocation->is_copy_allocation()) { - prefetch_time = static_cast(allocation) - ->copy_start_schedule_after(); - } else if (allocation->is_sliced_copy_allocation()) { - prefetch_time = static_cast(allocation) - ->slice_details_sorted_by_start_time() - .front() - .copy_start_after_time; - } else { - LOG(FATAL) << "Prefetch allocation are expected to be " - "CopyAllocations or SlicedCopyAllocations."; - } - if (prefetch_time != *request.preferred_prefetch_time) { - VLOG(1) << "Scheduled prefetch time (" << prefetch_time - << ") doesn't match the preferred prefetch time (" - << *request.preferred_prefetch_time - << "): " << request.use->hlo_use.ToString(); - if (request.strict_timing) { - return AllocationResult::kFailLiveRangeTooShort; - } - } - } - return AllocationResult::kSuccess; - } - // Warn if there was a preferred prefetch time but we couldn't actually - // prefetch. - if (request.preferred_prefetch_time) { - VLOG(1) << "The request has a preferred prefetch time (" - << *request.preferred_prefetch_time - << ") which could not be satisfied: " - << request.use->hlo_use.ToString(); - } - if (request.fail_on_unsatisfied_override || request.strict_timing) { - return prefetch_result; - } - result_mark(prefetch_result, allocation_result); - } - - CHECK(!request.require_end_colored_in_alternate_memory) - << "Failed to allocate end in alternate memory, even though its " - "required. " - "requires_contiguous_allocation: " - << request.allocation_value->requires_contiguous_allocation() - << " only_extend_existing_allocation: " - << request.only_extend_existing_allocation - << " require_end_colored_in_default_memory: " - << request.require_end_colored_in_default_memory - << " required_memory_space_at_end: " - << (required_memory_space_at_end == MemorySpace::kDefault ? "default" - : "alternate"); - - // If the end assignment was required to be in alternate memory but that - // wasn't possible, then this allocation is invalid. - if (required_memory_space_at_end == MemorySpace::kAlternate) { - return result_mark(AllocationResult::kFailRequiresUncommit, - allocation_result); - } - - // If the start assignment was required to be in alternate memory and the - // buffer needs a contiguous assignment, we couldn't satisfy this requirement - // and must abort. - if (required_memory_space_at_start == MemorySpace::kAlternate && - request.allocation_value->requires_contiguous_allocation()) { - return result_mark(AllocationResult::kFailRequiresUncommit, - allocation_result); - } - - // If a copy wasn't inserted, then add this use to the latest allocation in - // default memory. - (*prev_allocation_in_default_mem_it)->Extend(request.end_time); - (*prev_allocation_in_default_mem_it)->AddUse(request.use->hlo_use); - if (uses_in_default_memory_set_.insert(request.use->hlo_use).second) { - uses_in_default_memory_.push_back(request.use->hlo_use); - } - return allocation_result; -} - -void MsaAlgorithm::AddAsyncCopyForWindowPrefetch( - Allocation& prev_allocation, HloUse use, const Chunk& chunk, - int64_t exclusive_start_time, int64_t inclusive_end_time, - AllocationSequence* allocations, AliasedOffset* aliased_offset, - float resource, const WindowPrefetchedAllocation::Options& options) { - allocations->push_back(std::make_unique( - prev_allocation, use, chunk, exclusive_start_time, inclusive_end_time, - options)); - - RegisterAsyncCopy(MemorySpace::kAlternate, exclusive_start_time, - inclusive_end_time, allocations, aliased_offset, resource, - /*cross_program_prefetch_index=*/std::nullopt); -} - -void MsaAlgorithm::AddAsyncCopyOrOtherMemOp( - Allocation& prev_allocation, MemorySpace memory_space, - std::optional chunk, int64_t exclusive_start_time, int64_t end_time, - int64_t copy_done_schedule_before_time, AllocationSequence* allocations, - AliasedOffset* aliased_offset, float resource, - std::optional cross_program_prefetch_index, - HloInstruction* sync_mem_op, HloInstruction* async_mem_op_start, - HloInstruction* async_mem_op_done, int64_t source_operand_index) { - VLOG(3) << "Copy to " - << (memory_space == MemorySpace::kDefault ? "default" : "alternate") - << " memory in (" << exclusive_start_time << ", " - << copy_done_schedule_before_time << "), keeping until " << end_time - << ", estimated copy resource is " << resource; - CHECK_LT(exclusive_start_time, copy_done_schedule_before_time); - - allocations->push_back(std::make_unique( - prev_allocation, memory_space, chunk, exclusive_start_time, - copy_done_schedule_before_time, end_time, cross_program_prefetch_index, - sync_mem_op, async_mem_op_start, async_mem_op_done, - source_operand_index)); - - RegisterAsyncCopy(memory_space, exclusive_start_time, - copy_done_schedule_before_time, allocations, aliased_offset, - resource, cross_program_prefetch_index); -} - -void MsaAlgorithm::RegisterAsyncCopy( - MemorySpace memory_space, int64_t exclusive_start_time, - int64_t copy_done_schedule_before_time, AllocationSequence* allocations, - AliasedOffset* aliased_offset, float resource, - std::optional cross_program_prefetch_index) { - // Register the additional async copy with the interval tree to keep track of - // the limit at any given time. - pending_async_copies_.push_back({exclusive_start_time, - copy_done_schedule_before_time, resource, - memory_space, next_async_copy_id_++}); - if (memory_space == MemorySpace::kAlternate) { - prefetch_interval_tree_.Add( - /*start=*/ - ExclusiveToInclusiveStartTime(exclusive_start_time), - copy_done_schedule_before_time, kDummyChunk); - prefetch_async_copy_resource_.AddCopy(pending_async_copies_.back()); - if (options_.enforce_prefetch_fifo_order) { - async_copy_ordering_.AddCopy(pending_async_copies_.back()); - } - MaybeCreateOrAddToAliasedOffset(*allocations->back(), aliased_offset); - } else { - eviction_interval_tree_.Add( - /*start=*/ - ExclusiveToInclusiveStartTime(exclusive_start_time), - copy_done_schedule_before_time, kDummyChunk); - eviction_async_copy_resource_.AddCopy(pending_async_copies_.back()); - } -} - -namespace { - -// Computes a string that can be used for logging/debugging. For each slice, the -// string includes: -// - When the slice starts -// - When the slice copy must complete -// - When the allocation for the slice ends -// - An estimation of how much copy resource the slice consumes -std::string SliceTimesAndCopyResourcesToString( - const std::vector& slice_decisions, int64_t prefetch_end, - int64_t allocation_end) { - std::vector slice_strings; - slice_strings.reserve(slice_decisions.size()); - - for (const auto& slice_decision : slice_decisions) { - std::vector details; - details.push_back(absl::StrCat(slice_decision.exclusive_start_time)); - details.push_back(absl::StrCat(prefetch_end)); - details.push_back(absl::StrCat(allocation_end)); - details.push_back(absl::StrCat(slice_decision.copy_resource_consumed)); - - slice_strings.push_back( - absl::StrCat("(", absl::StrJoin(details, ", "), ")")); - } - - return absl::StrCat( - "Slices(copy_start_time, copy_done_by_time, allocation_end, " - "estimated_copy_resource) = [", - absl::StrJoin(slice_strings, ", "), "]"); -} - -} // namespace - -void MsaAlgorithm::AddAsyncSlicesForPrefetch( - const Allocation& prev_allocation, AllocationSequence* allocations, - AliasedOffset* aliased_offset, - const std::vector& slice_decisions_sorted_by_start_time, - int64_t prefetch_end_time, int64_t allocation_end_time, - HloInstruction* sync_mem_op) { - VLOG(3) << "Sliced copy to alternate memory. " - << SliceTimesAndCopyResourcesToString( - slice_decisions_sorted_by_start_time, prefetch_end_time, - allocation_end_time); - CHECK(absl::c_all_of( - slice_decisions_sorted_by_start_time, [&](const auto& slice_decision) { - return slice_decision.exclusive_start_time < prefetch_end_time; - })); - - allocations->push_back(std::make_unique( - prev_allocation, MemorySpace::kAlternate, - slice_decisions_sorted_by_start_time, prefetch_end_time, - allocation_end_time, options_.sliced_prefetch_options, - options_.get_equivalent_s8_shape_fn, sync_mem_op)); - - // Register the additional async copy with the interval tree to keep track of - // the limit at any given time. - for (const auto& slice_decision : slice_decisions_sorted_by_start_time) { - pending_async_copies_.push_back( - {slice_decision.exclusive_start_time, prefetch_end_time, - slice_decision.copy_resource_consumed, MemorySpace::kAlternate, - next_async_copy_id_++}); - prefetch_interval_tree_.Add(slice_decision.exclusive_start_time, - prefetch_end_time, kDummyChunk); - prefetch_async_copy_resource_.AddCopy(pending_async_copies_.back()); - if (options_.enforce_prefetch_fifo_order) { - async_copy_ordering_.AddCopy(pending_async_copies_.back()); - } - } - MaybeCreateOrAddToAliasedOffset(*allocations->back(), aliased_offset); -} - -bool MsaAlgorithm::ViolatesMaximumOutstandingAsyncCopies( - int64_t inclusive_start_time, int64_t end_time, bool is_prefetch, - int64_t extra_async_copy_limit, int64_t num_additional_copies) const { - if (options_.max_outstanding_prefetches < 0 && is_prefetch) { - return false; - } - if (options_.max_outstanding_evictions < 0 && !is_prefetch) { - return false; - } - - // Count the prefetches/evictions in the interval tree for the given interval. - if (is_prefetch) { - int64_t num_prefetches = prefetch_interval_tree_.NumChunksOverlappingInTime( - inclusive_start_time, end_time) + - num_additional_copies; - return num_prefetches >= - options_.max_outstanding_prefetches + extra_async_copy_limit; - } - int64_t num_evictions = eviction_interval_tree_.NumChunksOverlappingInTime( - inclusive_start_time, end_time) + - num_additional_copies; - return num_evictions >= - options_.max_outstanding_evictions + extra_async_copy_limit; -} - -AllocationResult MsaAlgorithm::ForceAlternateMemoryAllocationForMinTime( - const AllocationRequest& request) { - CHECK_EQ(request.allocation_value->position(), - request.allocation_value_to_update->position()); - - MsaBufferInterval alternate_mem_interval = MsaBufferInterval{ - /*buffer=*/request.allocation_value->value(), - /*size=*/request.size, - /*start=*/request.inclusive_start_time, - /*end=*/request.inclusive_start_time, - /*colocations=*/{}, - /*need_allocation=*/true, - }; - - Chunk chunk_candidate = FindChunkCandidate(alternate_mem_interval); - - if (chunk_candidate.chunk_end() > options_.max_size_in_bytes) { - return AllocationResult::kFailOutOfMemory; - } - - AddToPendingChunks(alternate_mem_interval, chunk_candidate); - - const HloPosition& defining_position = request.allocation_value->position(); - request.allocation_value->mutable_allocation_sequence()->push_back( - std::make_unique( - defining_position, MemorySpace::kAlternate, chunk_candidate, - alternate_mem_interval.start, alternate_mem_interval.end)); - // Since we did not use request.preferred_offset, we pass nullptr to - // CreateOrAddToAliasedOffset. - MaybeCreateOrAddToAliasedOffset( - *request.allocation_value->allocation_sequence()->back(), - /*aliased_offset=*/nullptr); - return AllocationResult::kSuccess; -} - -AllocationResult MsaAlgorithm::AllocateInAlternateMemoryNoCopy( - const AllocationRequest& request) { - Allocation* prev_allocation = nullptr; - bool can_eliminate_copy = false; - if (request.allocation_value->allocation_sequence()->empty()) { - // There hasn't been any allocations for this interval so far. We can - // eliminate copy if the value can be placed in the alternate memory or is - // pinned. - can_eliminate_copy = request.require_no_copy_alternate_mem_allocation || - options_.is_allowed_in_alternate_mem_fn( - *request.allocation_value->value()); - } else { - // If there has been a previous allocation, we can eliminate the copy if the - // previous allocation was also in the alternate memory. - prev_allocation = - request.allocation_value->allocation_sequence()->back().get(); - can_eliminate_copy = - (prev_allocation->memory_space() == MemorySpace::kAlternate && - !prev_allocation->is_window_prefetched_allocation()); - } - - if (!can_eliminate_copy && - (request.require_start_colored_in_alternate_memory || - request.require_end_colored_in_alternate_memory) && - (prev_allocation == nullptr || - prev_allocation->memory_space() == MemorySpace::kAlternate)) { - // If it is explicitly colored in alternate memory and the previous - // allocation is already in alternate memory (or nullptr), we can place it - // there. - can_eliminate_copy = true; - } - - if (!can_eliminate_copy) { - VLOG(3) << "Can't eliminate copy."; - return AllocationResult::kFailPrevAllocationNotInAlternateMem; - } - - const HloPosition& defining_position = request.allocation_value->position(); - // If prefer_no_copy_alternate_mem_allocation is true, bypass the live range - // duration checks. - if (!request.require_no_copy_alternate_mem_allocation && - !request.prefer_no_copy_alternate_mem_allocation && - !request.require_start_colored_in_alternate_memory && - !request.require_end_colored_in_alternate_memory && - !options_.prefetch_interval_picker->CanAllocateInAlternateMemoryNoCopy( - options_.cost_analysis ? options_.cost_analysis->GetShapeSizeBytes( - defining_position.shape()) - : 0, - request.inclusive_start_time, request.end_time)) { - VLOG(3) << "Live range is too long."; - return AllocationResult::kFailLiveRangeTooLong; - } - - MsaBufferInterval alternate_mem_interval; - alternate_mem_interval.buffer = request.allocation_value->value(); - alternate_mem_interval.size = request.size; - alternate_mem_interval.end = request.end_time; - alternate_mem_interval.start = request.inclusive_start_time; - - // Prefer the offset that was previously used for the previous allocation. - AliasedOffset* preferred_offset = nullptr; - if (prev_allocation != nullptr) { - preferred_offset = GetAliasedOffset(*prev_allocation); - // If there is a previous allocation, set the start time one after the end - // of the previous allocation's end. - alternate_mem_interval.start = prev_allocation->end_time() + 1; - if (request.no_copy_chunk_inclusive_start_time.has_value()) { - alternate_mem_interval.start = - *request.no_copy_chunk_inclusive_start_time; - } - } - - 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 satisfied. - if (preferred_offset && request.preferred_offset != preferred_offset) { - VLOG(3) << "Cannot perform no-copy allocation due to mismatch: " - "preferred_offset = " - << preferred_offset->offset << ", request.preferred_offset = " - << request.preferred_offset->offset; - return AllocationResult::kFailConflictingPreferredOffsets; - } - preferred_offset = request.preferred_offset; - } - - VLOG(3) << "We can eliminate copy to alternate memory. Preferred offset = " - << (preferred_offset ? preferred_offset->offset : -1); - // In case there are additional uses after this use, we rely on the last use - // time to try to reserve a chunk in the heap simulator. This is to prevent - // the following scenario: - // - // +-------+ - // / \ - // Producer--->Use1 +-->Use2 - // +---------+---------+ - // New buffer: | | | - // +---------+---------+ - // - // +-----------+ - // Current heap: | offset: 0 | - // --------------------------+-----------+------ - // - // Because we allocate buffers greedily, Producer to Use1 segment first, and - // then Use1 to Use2 segment, it is possible to allocate the first segment at - // an offset that is available for the first segment (e.g. offset 0) but not - // for the entire live range. This can result in unnecessary copies. By using - // the last use time, we try to find an allocation that is available for the - // entire Producer to Use2 range. - std::optional chunk_candidate = FindBestChunkCandidate( - request, preferred_offset, &alternate_mem_interval); - // Check if the new heap size fits within limits. Also ensure if a - // preferred offset was provided, that offset was used. - if (chunk_candidate) { - VLOG(3) << "Keep the buffer in alternate memory. Offset = " - << chunk_candidate->offset << ", size = " << chunk_candidate->size - << ", heap_size = " << result_.UpdatedHeapSize(*chunk_candidate) - << ", prefetch picker = " - << options_.prefetch_interval_picker->ToNoCopyDebugString( - defining_position.shape(), - /*start_time=*/ - InclusiveToExclusiveStartTime(request.inclusive_start_time), - request.end_time); - AddToPendingChunks(alternate_mem_interval, *chunk_candidate); - - // If there was a previous allocation, the buffer location is the - // same as the previous. Otherwise, it is the operand. - if (prev_allocation != nullptr && - (prev_allocation->is_copy_like_allocation() || - prev_allocation->defining_position() == defining_position)) { - prev_allocation->Extend(request.end_time); - } else { - request.allocation_value->mutable_allocation_sequence()->push_back( - std::make_unique( - defining_position, MemorySpace::kAlternate, chunk_candidate, - request.inclusive_start_time, request.end_time)); - MaybeCreateOrAddToAliasedOffset( - *request.allocation_value->allocation_sequence()->back(), - preferred_offset); - } - if (!request.only_extend_existing_allocation) { - request.allocation_value->allocation_sequence()->back()->AddUse( - request.use->hlo_use); - } - return AllocationResult::kSuccess; - } - if (request.prefer_no_copy_alternate_mem_allocation) { - VLOG(1) << "Preferred no-copy allocation, but this was not possible: " - << request.use->hlo_use.ToString(); - } - return AllocationResult::kFailOutOfMemory; -} - -AllocationResult MsaAlgorithm::Evict(const AllocationRequest& request, - bool force_evict) { - CHECK_GT(request.allocation_value->allocation_sequence()->size(), 0); - Allocation* prev_allocation = - request.allocation_value->allocation_sequence()->back().get(); - // We do not ever expect an Evict() to be immediately proceeded by a prefetch. - // If that case ever occurs, the eviction_exclusive_start_time below will be - // calculated incorrectly, as it will need to come after the prefetch finishes - // coping data. - CHECK(!prev_allocation->is_copy_like_allocation()) - << "Evict has been given copy-like previous allocation.\nEvict " - "candidate:\n" - << request.allocation_value->ToString() << "\nPrevious allocation:\n" - << prev_allocation->ToString(); - - // The previous allocation's inclusive start time is the eviction's exclusive - // start time to ensure that the value is created before we start copying - // back to default memory. - int64_t eviction_exclusive_start_time = prev_allocation->start_time(); - // TODO(b/483747673): `prev_allocation->end_time()` is conservative for - // eviction done schedule before time. We can set it to - // `prev_allocation->end_time() + 1`. - int64_t eviction_end_time = prev_allocation->end_time(); - CHECK(eviction_exclusive_start_time <= eviction_end_time); - - int64_t preferred_eviction_end_time = - std::max(options_.prefetch_interval_picker->PreferredEvictionEndTime( - request.allocation_value_to_update->position().shape(), - eviction_exclusive_start_time, request.end_time), - eviction_end_time); - // Evictions must complete by the time of this use. - preferred_eviction_end_time = - std::min(preferred_eviction_end_time, request.latest_prefetch_time); - - MsaBufferInterval eviction_mem_interval; - eviction_mem_interval.buffer = request.allocation_value->value(); - // When replacing an sync slice, the size of the original allocation_value - // matters instead of the queuing_allocation_value - // TODO(mehrdadk): separate the request size for src and dst - // AllocationSequence - eviction_mem_interval.size = - std::max(request.allocation_value->size(), request.size); - // Try to reserve a buffer from the end of the previous allocation to the - // preferred eviction end time. - eviction_mem_interval.start = eviction_end_time + 1; - eviction_mem_interval.end = preferred_eviction_end_time; - int64_t preferred_offset = prev_allocation->chunk().offset; - VLOG(3) << "Considering eviction after" << eviction_exclusive_start_time - << ", with preferred end time = " << eviction_mem_interval.end; - - auto next_eviction_end_time_candidate = [&]() { - // Only call FindChunkCandidate if the index is an edge time (where there is - // another allocation that starts or ends at this time index). This reduces - // redundant calls to FundChunkCandidate and so reduces compilation time. - int64_t next_eviction_end_time = eviction_mem_interval.end; - do { - --next_eviction_end_time; - } while (next_eviction_end_time > eviction_end_time && - !edge_time_indices_.contains(next_eviction_end_time)); - return next_eviction_end_time; - }; - int64_t original_end_time = eviction_mem_interval.end; - std::optional final_chunk_candidate; - for (; eviction_mem_interval.end > eviction_end_time; - eviction_mem_interval.end = next_eviction_end_time_candidate()) { - Chunk chunk_candidate; - // If the buffer has no colocations, then use the fast algorithm to find - // the earliest end time with a free chunk at the preferred offset. - if (GetTransitiveColocations(eviction_mem_interval).empty()) { - int64_t earliest_end_with_free_chunk = - FindLatestEndWithFreeChunkAtPreferredOffset(eviction_mem_interval, - preferred_offset); - if (earliest_end_with_free_chunk <= eviction_end_time) { - eviction_mem_interval.end = eviction_end_time; - break; - } - eviction_mem_interval.end = earliest_end_with_free_chunk; - chunk_candidate = - Chunk::FromOffsetSize(preferred_offset, eviction_mem_interval.size); - } else { - chunk_candidate = - FindChunkCandidate(eviction_mem_interval, preferred_offset); - } - - if (chunk_candidate.offset == preferred_offset && - (eviction_mem_interval.end == original_end_time || - edge_time_indices_.contains(eviction_mem_interval.end))) { - final_chunk_candidate = chunk_candidate; - break; - } - } - eviction_end_time = eviction_mem_interval.end; - - VLOG(3) << "Evicting buffer at " << prev_allocation->chunk().offset << " (" - << eviction_exclusive_start_time << ", " << eviction_end_time << ")"; - - float eviction_resource = - options_.cost_analysis - ? options_.cost_analysis->GetAsyncCopyElapsed( - options_.cost_analysis->GetShapeSizeBytes( - request.allocation_value->position().shape())) - : 0.1; - - bool eviction_interval_too_short = - (eviction_exclusive_start_time == eviction_end_time); - bool eviction_violates_resource = - !eviction_async_copy_resource_.HasEnoughResource( - eviction_exclusive_start_time, eviction_end_time, eviction_resource); - if (eviction_violates_resource) { - // If we're in the last retry, set resource to 0. - if (options_.prefetch_interval_picker->retry_number() == - options_.max_retries - 1) { - VLOG(3) << "Violates resource in last retry, setting resource = 0"; - eviction_resource = 0; - } - eviction_violates_resource = - !eviction_async_copy_resource_.HasEnoughResource( - eviction_exclusive_start_time, eviction_end_time, - eviction_resource); - } - bool eviction_violates_outstanding_copies = - ViolatesMaximumOutstandingAsyncCopies( - /*inclusive_start_time=*/ExclusiveToInclusiveStartTime( - eviction_exclusive_start_time), - eviction_end_time, - /*is_prefetch=*/false); - - // See if this interval would violate the asynchronous copy limit. - if (!eviction_interval_too_short && !eviction_violates_outstanding_copies && - !eviction_violates_resource) { - prev_allocation->Extend(eviction_end_time); - if (final_chunk_candidate.has_value()) { - AddToPendingChunks(eviction_mem_interval, final_chunk_candidate.value()); - } - AddAsyncCopyOrOtherMemOp( - *prev_allocation, MemorySpace::kDefault, - /*chunk=*/std::nullopt, eviction_exclusive_start_time, - prev_allocation->end_time(), eviction_end_time, - request.allocation_value->mutable_allocation_sequence(), - /*aliased_offset=*/nullptr, eviction_resource); - return AllocationResult::kSuccess; - } - - if (eviction_violates_outstanding_copies) { - VLOG(3) << "This violates the maximum async copies."; - } else if (eviction_violates_resource) { - VLOG(3) << "This violates resource."; - } else { - VLOG(3) << "Eviction interval is too short (" - << eviction_exclusive_start_time << ", " << eviction_end_time - << ")."; - } - - if (force_evict) { - VLOG(3) << "Forcing evicting."; - int64_t inclusive_eviction_start_time = prev_allocation->end_time() + 1; - AddAsyncCopyOrOtherMemOp( - *prev_allocation, MemorySpace::kDefault, - /*chunk=*/std::nullopt, - /*exclusive_start_time=*/ - InclusiveToExclusiveStartTime(inclusive_eviction_start_time), - request.end_time, - /*copy_done_schedule_before_time=*/prev_allocation->end_time() + 1, - request.allocation_value->mutable_allocation_sequence(), - /*aliased_offset=*/nullptr, 0); - return AllocationResult::kSuccess; - } - - // If the eviction couldn't be scheduled, then fail. This buffer will be - // kept in the default memory. - VLOG(3) << "Bailing: Could not evict " << request.use->hlo_use.ToString() - << " because we hit the limit of maximum asynchronous copies " - << "between (" - << hlo_live_range_.flattened_instruction_sequence() - .instructions()[eviction_exclusive_start_time] - << ", " - << hlo_live_range_.flattened_instruction_sequence() - .instructions()[eviction_end_time] - << ")"; - return AllocationResult::kFailOutOfAsyncCopies; -} - -int64_t MsaAlgorithm::FindPrefetchEndTime( - const AllocationRequest& request, int64_t earliest_prefetch_time) const { - return request.latest_prefetch_time; -} - -namespace { - -// A debugging/logging method for describing a sliced solution. -std::string DescribeSlicedBufferMove( - const std::vector& slice_decisions, - const MsaAlgorithm::HeapResult& heap_result, - const MsaAlgorithm::Chunk& full_chunk, - absl::string_view prefetch_picker_debug_string) { - std::vector slice_strings; - slice_strings.reserve(slice_decisions.size()); - - for (const auto& slice_decision : slice_decisions) { - slice_strings.push_back(absl::StrCat( - "(", slice_decision.exclusive_start_time, ", ", - slice_decision.chunk.offset, ", ", slice_decision.chunk.size, ")")); - } - - return absl::StrCat( - "Moving buffer to alternate memory in slices. Slices(start_time, offset, " - "size) = [", - absl::StrJoin(slice_strings, ", "), - "]. Heap size = ", heap_result.UpdatedHeapSize(full_chunk), - ". Prefetch picker = ", prefetch_picker_debug_string); -} - -} // namespace - -void MsaAlgorithm::WindowPrefetchOperand(const HloUse& use, int64_t bytes) { - CHECK(options_.IsOpSpanExposureEnabled()); - - HloInstruction* instruction = use.instruction; - ShapeIndex shape_index = use.operand_index; - HloInstruction* operand = instruction->mutable_operand(use.operand_number); - // Find the defining position of the operand. - for (int i = 0; i < use.operand_index.size(); ++i) { - CHECK(operand->opcode() == HloOpcode::kGetTupleElement); - operand = operand->mutable_operand(0); - } - - // Create a new HloValue for the window buffer. - HloValue::Id new_value_id = alias_analysis_.dataflow_analysis().NewValueId(); - HloValue hlo_value(new_value_id, operand, shape_index); - int64_t start_time = hlo_live_range_.instruction_schedule().at(operand); - int64_t end_time = hlo_live_range_.instruction_schedule().at(instruction); - - // Create a buffer interval, which has the same start and end time as the - // operand. The hlo value is the operand. - MsaBufferInterval buffer_interval; - buffer_interval.buffer = &hlo_value; - buffer_interval.size = bytes; - buffer_interval.start = start_time; - buffer_interval.end = end_time; - buffer_interval.need_allocation = true; - - // Create an allocation_values using the buffer interval. - std::vector allocation_values; - allocation_values.emplace_back(&hlo_value, hlo_value.defining_position(), - bytes); - allocation_values[0].AddUse(use, end_time); - - // Create an allocation request using the allocation_value. - AllocationValue& allocation_value = allocation_values[0]; - AllocationValue::Use& allocation_value_use = allocation_value.uses()[0]; - std::vector all_use_times = {end_time}; - AllocationRequest request = CreateAllocationRequest( - allocation_value, allocation_value, - /*use=*/allocation_value_use, /*previous_use=*/nullptr, - /*preferred_offset=*/nullptr, - /*definition_time=*/start_time, - /*require_no_copy_alternate_mem_allocation=*/false, - /*all_use_times=*/all_use_times, - /*only_extend_existing_allocation=*/false, - /*processed_allocation_values=*/{}, - /*all_allocation_values=*/{}, - /*shape_override=*/ - ShapeUtil::MakeValidatedShape(U8, {bytes}).value()); - - // Create a dummy allocation that is in the default memory, this is needed for - // creating a WindowPrefetchedAllocation. This allocation does not need to be - // appended to the allocation sequence. - PinnedAllocation dummy_prev_allocation( - /*defining_position=*/{operand, shape_index}, MemorySpace::kDefault, - /*chunk=*/std::nullopt, request.inclusive_start_time, request.end_time); - - // Construct the options needed for creating the window prefetch allocation. - WindowPrefetchedAllocation::Options options; - options.bytes = bytes; - options.alternate_memory_space = options_.alternate_memory_space; - options.notify_operand_appended_fn = options_.notify_operand_appended_fn; - request.window_prefetch_options = &options; - - if (options_.IsWindowPrefetchingEnabled()) { - // Window prefetch mode - Prefetch(request, dummy_prev_allocation); - } else { - // Window exposure mode, we only need to find a chunk for the window - // buffer. - CHECK(options_.IsOpSpanExposureEnabled() && - !options_.IsWindowPrefetchingEnabled()); - // Adjust the start time of the buffer interval to be the use time. This is - // because we only need the buffer to be alive at the use time. - buffer_interval.start = end_time; - std::optional candidate_chunk = FindBestChunkCandidate( - request, /*preferred_offset=*/nullptr, &buffer_interval); - if (candidate_chunk.has_value()) { - AddToPendingChunks(buffer_interval, *candidate_chunk); - - AllocationSequence* allocation_sequence = - allocation_value.mutable_allocation_sequence(); - allocation_sequence->push_back( - std::make_unique( - dummy_prev_allocation, use, *candidate_chunk, end_time - 1, - end_time, options)); - MaybeCreateOrAddToAliasedOffset(*allocation_sequence->back(), - /*aliased_offset=*/nullptr); - allocation_sequence->back()->AddUse(use); - } - } - // Finalize the allocation values. This adds the newly created allocation to - // allocations_. - FinalizeAllocations(absl::MakeSpan(allocation_values)); -} - -absl::Status MsaAlgorithm::WindowPrefetch() { - CHECK(options_.IsOpSpanExposureEnabled()); - - absl::flat_hash_set window_prefetchable_instructions; - - // At this point, we don't have the memory space colored in the original - // instruction, but we know which operands and outputs are in the alternate - // memory. So we clone the instruction and color the operands and outputs that - // are in the alternate memory. Then we propagate the memory space to the - // cloned computation and use the cloned computation to determine the operand - // span size. - - // Map of the original instruction to a clone of the instruction. Use a - // vector to ensure deterministic traversal for memory space propagation - // and cleanup. - std::vector cloned_insts_order; - absl::flat_hash_map cloned_insts; - const std::vector& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - - for (HloInstruction* instruction : instruction_sequence) { - if (!options_.is_window_prefetchable_instruction_fn(instruction)) { - continue; - } - - window_prefetchable_instructions.insert(instruction); - - // This lambda sets an hlo's memory space to the alternate memory space. - auto color_hlo = [&](HloInstruction* hlo, ShapeIndex shape_index) { - ShapeUtil::GetMutableSubshape(hlo->mutable_shape(), shape_index) - ->mutable_layout() - ->set_memory_space(options_.alternate_memory_space); - }; - - // Make a clone of the instruction. - HloInstruction* cloned = - instruction->parent()->AddInstruction(instruction->Clone()); - cloned_insts[instruction] = cloned; - cloned_insts_order.push_back(instruction); - - // Color the cloned instruction's fused parameters. - auto it = operands_in_alternate_memory_map_.find(instruction); - if (it != operands_in_alternate_memory_map_.end()) { - for (const auto& [i, _] : it->second) { - // For the operand in the alternate memory, color the parameter. Because - // we color the fused parameters, it is the full shape of the parameter. - // By using the full shape, we assume that the parameter is not a - // tuple and one or more of the tensors in that parameter are in - // alternate memory. Let's add a check to make sure that the parameter - // is not a tuple. - if (cloned->fused_parameters()[i]->shape().IsTuple()) { - LOG(ERROR) << "Tuple parameter: " - << cloned->fused_parameters()[i]->shape().ToString(true); - return absl::FailedPreconditionError( - "Tuple parameter not supported for window prefetch."); - } - color_hlo(cloned->fused_parameters()[i], {}); - } - } - - // Color the cloned instruction's outputs. - if (auto it = outputs_in_alternate_memory_map_.find(instruction); - it != outputs_in_alternate_memory_map_.end()) { - for (const auto& shape_index : it->second) { - color_hlo(cloned->fused_expression_root(), shape_index); - } - } - } - - // Propagate the memory space to the cloned fusion computations. - ABSL_ASSIGN_OR_RETURN(auto dataflow_analysis, - HloDataflowAnalysis::Run(*module_, /*ssa_form=*/false, - /*bitcast_defines_value=*/true)); - MemorySpacePropagation memory_space_propagation(std::move(dataflow_analysis)); - for (HloInstruction* inst : cloned_insts_order) { - HloInstruction* cloned = cloned_insts[inst]; - for (HloComputation* computation : cloned->called_computations()) { - memory_space_propagation.RunOnComputation(computation); - } - } - - int64_t window_prefetch_operand_count = 0; - // Prefetch the window buffers. - for (const HloUse& use : uses_in_default_memory_) { - if (!window_prefetchable_instructions.contains(use.instruction)) { - continue; - } - - CHECK(options_.op_span_size_fn); - int64_t span_size = options_.op_span_size_fn( - use.instruction, cloned_insts[use.instruction], use.operand_number); - if (span_size >= options_.window_prefetch_min_span_size && - window_prefetch_operand_count < options_.window_prefetch_max_operands) { - window_prefetch_operand_count++; - WindowPrefetchOperand(use, span_size); - } - } - - // Remove the cloned instructions. - for (HloInstruction* inst : cloned_insts_order) { - HloInstruction* cloned = cloned_insts[inst]; - HloComputation* computation = cloned->parent(); - CHECK_OK(computation->RemoveInstruction(cloned)); - computation->Cleanup(); - } - return absl::OkStatus(); -} - -AllocationResult MsaAlgorithm::Prefetch( - const AllocationRequest& request, - Allocation& prev_allocation_in_default_mem, bool force_prefetch) { - AllocationResult result = - PrefetchWithResourceConstraints(request, prev_allocation_in_default_mem); - if (result == AllocationResult::kSuccess || !force_prefetch) { - return result; - } - - // We require the buffer in alternate memory for the use, but, prefetching - // failed, so we force a prefetch and point the use to the prefetch - // allocation. - MsaBufferInterval alternate_mem_interval = MsaBufferInterval{ - /*buffer=*/request.allocation_value->value(), - /*size=*/request.size, - /*start=*/request.end_time, - /*end=*/request.end_time, - /*colocations=*/{}, - /*need_allocation=*/true, - }; - - if (options_.enforce_prefetch_fifo_order) { - CHECK(!async_copy_ordering_.ViolatesOrdering( - /*exclusive_start_time=*/alternate_mem_interval.start - 1, - /*end_time=*/alternate_mem_interval.end)) - << "Scheduling immediate prefetch is necessary to satisfy coloring " - "constraints but it violates FIFO ordering. Consider coloring " - "lesser buffers in alternate memory or disable prefetch fifo " - "ordering."; - } - - Chunk chunk_candidate = FindChunkCandidate(alternate_mem_interval); - - if (chunk_candidate.chunk_end() > options_.max_size_in_bytes) { - return AllocationResult::kFailOutOfMemory; - } - AddToPendingChunks(alternate_mem_interval, chunk_candidate); - - AddAsyncCopyOrOtherMemOp( - prev_allocation_in_default_mem, MemorySpace::kAlternate, chunk_candidate, - request.end_time - 1, request.end_time, request.end_time, - request.allocation_value->mutable_allocation_sequence(), - /*aliased_offset=*/nullptr, 0.0f, - /*cross_program_prefetch_index=*/std::nullopt, - request.required_copy_allocation_for); - request.allocation_value->allocation_sequence()->back()->AddUse( - request.use->hlo_use); - return AllocationResult::kSuccess; -} - -AllocationResult MsaAlgorithm::PrefetchWithResourceConstraints( - const AllocationRequest& request, - Allocation& prev_allocation_in_default_mem) { - // Try partially placing the buffer in the alternate space. The time that is - // overlapped will be used to asynchronously copy the buffer from the - // default memory to the alternate memory. - // - // start end - // time time - // X---------------------X - // Alternate: +------+ - // Default: +---------------------+ - // ^ ^ - // Copy Copy - // Start Done - - VLOG(5) << "Considering prefetch of " - << request.allocation_value->ToShortString() - << (request.preferred_offset - ? absl::StrCat(", with a preferred offset of ", - request.preferred_offset->offset, ".") - : ""); - PrefetchContext context; - context.request = &request; - context.prev_allocation_in_default_mem = &prev_allocation_in_default_mem; - // If the request has window prefetch options, it is called from window - // prefetch. - context.window_prefetch = (request.window_prefetch_options != nullptr); - // If we're actually doing a prefetch, it should be for a non-window-prefetch - // buffer or the window prefetching should be enabled. - CHECK(!context.window_prefetch || options_.IsWindowPrefetchingEnabled()); - - // Create a SliceProposal and WorkingIntervals. - SetupPrefetchWorkingIntervalsAndSliceProposal(context); - - // Compute some additional preliminaries - AllocationResult init_result = InitializePrefetchIntervalPicker(context); - if (init_result != AllocationResult::kSuccess) { - return init_result; - } - AllocationResult check_result = EnsureSomeSpatialPrefetchFitExists(context); - if (check_result != AllocationResult::kSuccess) { - return check_result; - } - const HloUse& use = request.use->hlo_use; - if (request.shape_override.has_value()) { - context.full_shape = &*request.shape_override; - } else { - context.full_shape = &ShapeUtil::GetSubshape( - use.instruction->operand(use.operand_number)->shape(), - use.operand_index); - } - // While uses might be allowed to have additional outstanding prefetches. - context.extra_async_copy_limit = - use.instruction->opcode() == HloOpcode::kWhile - ? options_.while_use_extra_outstanding_prefetch_limit - : 0; - - // If the request is for a sync mem op conversion to async, we may allow for - // more async copies. - if (context.request->require_copy_allocation) { - context.extra_async_copy_limit += - options_.extend_async_copies_limit_for_sync_mem_op_conversion; - } - - // Loop over potential prefetch starting times. At the selected start time, we - // check if we have enough resources and memory for a sliced version of the - // request and a non-sliced version of the request. We return the first sliced - // solution that we find. We fallback to the first unsliced solution we find, - // if we are unable to find a sliced solution. - AllocationResult result = AllocationResult::kSuccess; - while (!options_.prefetch_interval_picker->Done()) { - // Get the prefetch start time from the interval picker. - context.exclusive_prefetch_start_time = - options_.prefetch_interval_picker->Next(); - CHECK_LT(context.exclusive_prefetch_start_time, context.prefetch_end_time); - if (context.exclusive_out_of_mem_start.has_value() && - context.exclusive_prefetch_start_time <= - *context.exclusive_out_of_mem_start) { - VLOG(4) << "This would OOM (cached)."; - return AllocationResult::kFailOutOfMemory; - } - - if (context.slice_proposal_collection) { - VLOG(5) << "Trying sliced solution."; - // Check if a sliced solution fits. - AllocationResult sliced_result = - CheckPrefetchFit(/*for_sliced_solution=*/true, context); - if (sliced_result == AllocationResult::kSuccess) { - // Break out of the loop and use the sliced solution. - CHECK(context.sliced_solution); - break; - } else if (sliced_result != - AllocationResult::kAllSlicesHaveTheSameStartTime) { - result_mark(sliced_result, result); - } - } - - // If we don't already have an unsliced solution, check the current fit. - if (!context.unsliced_solution) { - VLOG(5) << "Trying unsliced solution."; - AllocationResult unsliced_result = - CheckPrefetchFit(/*for_sliced_solution=*/false, context); - if (unsliced_result != AllocationResult::kSuccess) { - result_mark(unsliced_result, result); - } else if (!context.slice_proposal_collection) { - // We found an unsliced solution and there is no slice proposal, so - // break out of the loop and use the unsliced solution. - CHECK(context.unsliced_solution); - break; - } - } - } - - // Check if we found any solutions. - if (context.sliced_solution && !context.request->required_copy_for_slice) { - CHECK(!context.sliced_solution->slices_for_pending_chunks.empty()); - VLOG(3) << DescribeSlicedBufferMove( - context.sliced_solution->slice_decisions_sorted_by_start_time, result_, - context.sliced_solution->slices_for_pending_chunks.back().second, - context.sliced_solution->prefetch_picker_debug_string); - - for (const auto& interval_chunk_pair : - context.sliced_solution->slices_for_pending_chunks) { - AddToPendingChunks(interval_chunk_pair.first, interval_chunk_pair.second); - } - AddAsyncSlicesForPrefetch( - *context.prev_allocation_in_default_mem, - context.request->allocation_value_to_update - ->mutable_allocation_sequence(), - context.request->preferred_offset, - context.sliced_solution->slice_decisions_sorted_by_start_time, - context.prefetch_end_time, context.request->end_time, - context.request->required_copy_allocation_for); - context.request->allocation_value_to_update->allocation_sequence() - ->back() - ->AddUse(context.request->use->hlo_use); - return AllocationResult::kSuccess; - } - if (context.unsliced_solution) { - VLOG(3) << "Move the buffer to alternate memory after time " - << InclusiveToExclusiveStartTime( - context.unsliced_solution_intervals.full.start) - << ". Offset = " - << context.unsliced_solution->chunk_candidate.offset - << ", size = " << context.unsliced_solution->chunk_candidate.size - << ", heap_size = " - << result_.UpdatedHeapSize( - context.unsliced_solution->chunk_candidate) - << ", prefetch picker = " - << context.unsliced_solution->prefetch_picker_debug_string; - AddToPendingChunks(context.unsliced_solution_intervals.full, - context.unsliced_solution->chunk_candidate); - if (context.window_prefetch) { - AddAsyncCopyForWindowPrefetch( - *context.prev_allocation_in_default_mem, request.use->hlo_use, - context.unsliced_solution->chunk_candidate, - context.unsliced_solution_intervals.full.start - 1, - context.prefetch_end_time, - context.request->allocation_value_to_update - ->mutable_allocation_sequence(), - context.request->preferred_offset, - context.unsliced_solution->prefetch_resource, - *context.request->window_prefetch_options); - } else { - AddAsyncCopyOrOtherMemOp( - *context.prev_allocation_in_default_mem, MemorySpace::kAlternate, - context.unsliced_solution->chunk_candidate, - context.unsliced_solution_intervals.full.start - 1, - context.request->end_time, context.prefetch_end_time, - context.request->allocation_value_to_update - ->mutable_allocation_sequence(), - context.request->preferred_offset, - context.unsliced_solution->prefetch_resource, - /*cross_program_prefetch_index=*/std::nullopt, - context.request->required_copy_allocation_for); - context.prev_allocation_in_default_mem->Extend( - context.request->latest_prefetch_time); - } - - request.allocation_value_to_update->allocation_sequence()->back()->AddUse( - request.use->hlo_use); - return AllocationResult::kSuccess; - } - - // If we didn't consider any prefetch intervals, then the live range was too - // short. - return (result == AllocationResult::kSuccess - ? AllocationResult::kFailLiveRangeTooShort - : result); -} - -void MsaAlgorithm::GenerateSliceProposal(PrefetchContext& context) const { - if (context.request->required_copy_for_slice) { - VLOG(5) << "Not slicing " << context.request->use->hlo_use - << " because slicing a slice instruction is not supported yet."; - return; - } - - if (options_.sliced_prefetch_options.max_slices() < 2) { - return; - } - auto log_prefix = [&]() { - return absl::StrCat( - "Slice request(options = ", - options_.sliced_prefetch_options.ShortDebugString(), "; shape = ", - context.prev_allocation_in_default_mem->defining_position() - .shape() - .ToString(), - ")"); - }; - - if (context.request->size < options_.sliced_prefetch_options.min_bytes()) { - VLOG(5) << "Not slicing " << log_prefix() << " because the request size " - << context.request->size - << " is smaller than the min configured size of " - << options_.sliced_prefetch_options.min_bytes(); - return; - } - - auto status_or_proposal = options_.propose_slice_fn( - context.prev_allocation_in_default_mem->defining_position().shape(), - options_.sliced_prefetch_options); - if (!status_or_proposal.ok()) { - VLOG(2) << log_prefix() << " failed: " << status_or_proposal.status(); - return; - } - - if (status_or_proposal.value().size() < 2) { - VLOG(2) << log_prefix() << ". No slices proposed."; - return; - } - - VLOG(6) << log_prefix() << ". Slice proposal = [" - << absl::StrJoin(status_or_proposal.value(), ", ", - [](std::string* out, const SliceProposal& proposal) { - absl::StrAppend(out, proposal.ToString()); - }) - << "]"; - - context.slice_proposal_collection = std::move(status_or_proposal.value()); -} - -void MsaAlgorithm::SetupPrefetchWorkingIntervalsAndSliceProposal( - PrefetchContext& context) const { - // Setup the full WorkingIntervals for the sliced and unsliced solutions. - // Future code will adjust the start and end times. - context.sliced_solution_intervals.full = MsaBufferInterval{ - context.request->allocation_value_to_update->value(), - /*size=*/context.request->size, - /*start=*/-1, - /*end=*/context.request->end_time, - /*colocations=*/{}, - /*need_allocation=*/true, - }; - context.unsliced_solution_intervals.full = - context.sliced_solution_intervals.full; - - // Attempt to generate a slice proposal. - if (!context.window_prefetch) { - GenerateSliceProposal(context); - } - - // Setup the full SlicedBufferIntervals for the sliced and unsliced solutions. - // If there is no slice proposal, we will not try a sliced solution. In such a - // case, we do not populate context.sliced_solution_intervals. - if (context.slice_proposal_collection) { - context.sliced_solution_intervals.sliced = - std::make_unique( - SlicedBufferInterval::CreateMutableInterval( - context.sliced_solution_intervals.full)); - std::vector sizes; - sizes.reserve(context.slice_proposal_collection->size()); - for (const SliceProposal& single_slice_proposal : - *context.slice_proposal_collection) { - sizes.push_back(single_slice_proposal.slice_size); - } - context.sliced_solution_intervals.sliced->Slice(sizes); - } - context.unsliced_solution_intervals.sliced = - std::make_unique( - SlicedBufferInterval::CreateMutableInterval( - context.unsliced_solution_intervals.full)); -} - -AllocationResult MsaAlgorithm::InitializePrefetchIntervalPicker( - PrefetchContext& context) { - int64_t earliest_exclusive_prefetch_time = - context.prev_allocation_in_default_mem->earliest_available_time(); - if (context.request->earliest_prefetch_time) { - earliest_exclusive_prefetch_time = - std::max(earliest_exclusive_prefetch_time, - *context.request->earliest_prefetch_time); - } - context.prefetch_end_time = - FindPrefetchEndTime(*context.request, earliest_exclusive_prefetch_time); - - // As a compile time optimization, use the peak memory usage to filter out - // allocation times that would push us to OOM. - std::optional earliest_exclusive_non_oom_prefetch_time = - FindEarliestExclusiveTimeToSatisfyPeakMemory( - earliest_exclusive_prefetch_time, context.prefetch_end_time, - context.request->size); - if (!earliest_exclusive_non_oom_prefetch_time) { - VLOG(3) << "Any prefetch in range (" << earliest_exclusive_prefetch_time - << ", " << context.prefetch_end_time << ") for size " - << context.request->size << " would go out of memory."; - return AllocationResult::kFailOutOfMemory; - } - if (!context.slice_proposal_collection) { - // We can only perform this optimization if we are not slicing. - // earliest_non_oom_prefetch_time lets us know the first time the entire - // buffer will fit, but we may be able to start slices before that time. So, - // we leave earliest_prefetch_time at its initial value. - VLOG(4) << "After peak memory check, prefetch range is (" - << *earliest_exclusive_non_oom_prefetch_time << ", " - << context.prefetch_end_time - << "). Original earliest prefetch time is " - << earliest_exclusive_prefetch_time; - earliest_exclusive_prefetch_time = - *earliest_exclusive_non_oom_prefetch_time; - } - std::optional preferred_prefetch_time = - context.request->preferred_prefetch_time; - if (preferred_prefetch_time) { - if (context.request->strict_timing) { - if (*preferred_prefetch_time < earliest_exclusive_prefetch_time || - *preferred_prefetch_time >= context.prefetch_end_time) { - VLOG(3) << "Strict prefetch time " << *preferred_prefetch_time - << " is outside allowable prefetch range (" - << earliest_exclusive_prefetch_time << ", " - << context.prefetch_end_time << ")."; - return AllocationResult::kFailLiveRangeTooShort; - } - } else { - preferred_prefetch_time = - std::max(*preferred_prefetch_time, earliest_exclusive_prefetch_time); - } - } - options_.prefetch_interval_picker->Begin( - context.request->use->hlo_use, earliest_exclusive_prefetch_time, - context.prefetch_end_time, preferred_prefetch_time, - context.request->strict_timing); - VLOG(3) << "Trying prefetch picker = " - << options_.prefetch_interval_picker->ToDebugString(); - - return AllocationResult::kSuccess; -} - -AllocationResult MsaAlgorithm::EnsureSomeSpatialPrefetchFitExists( - PrefetchContext& context) const { - SlicedBufferInterval* interval = - (context.slice_proposal_collection - ? context.sliced_solution_intervals.sliced.get() - : context.unsliced_solution_intervals.sliced.get()); - - // Note, UpdateInclusiveSliceStartTimes() will correctly update start times - // for both sliced and unsliced solutions. - interval->UpdateExclusiveSliceStartTimes( - std::vector(interval->num_slices(), - options_.prefetch_interval_picker->latest_time())); - std::vector chunk_candidates = FindBestChunkCandidates( - *context.request, context.request->preferred_offset, interval); - if (chunk_candidates.empty()) { - VLOG(3) << "The latest prefetch (" << interval->full_buffer_interval().start - << ", " << context.request->end_time - << ") cannot find valid chunks. Giving up."; - return AllocationResult::kFailOutOfMemory; - } - - return AllocationResult::kSuccess; -} - -namespace { - -// Returns the copy resources needed for the specified slice proposal -// collection, in descending order. -std::vector GetCopyResourcesSortedDescending( - const Options& options, - const SliceProposalCollection& slice_proposal_collection) { - std::vector copy_resources; - copy_resources.reserve(slice_proposal_collection.size()); - for (const SliceProposal& proposal : slice_proposal_collection) { - copy_resources.push_back( - CopyResourceForShape(options, proposal.slice_shape)); - } - absl::c_sort(copy_resources); - return copy_resources; -} - -// Returns true if we would have enough async copy resources to copy each -// specified slice. -bool DoWeHaveEnoughCopyResource( - const std::vector& slice_start_times, int64_t prefetch_end_time, - const std::vector& copy_resource_per_slice, - AsynchronousCopyResource& async_copy_resource) { - CHECK_EQ(slice_start_times.size(), copy_resource_per_slice.size()); - - std::vector specs; - specs.reserve(slice_start_times.size()); - - // Note, the HasEnoughResourceMultiCheck() below is sensitive to this order. - // The specs must be in slice start time order because that's the order - // they'll be added to prefetch_async_copy_resource_ in - // AddAsyncSlicesForPrefetch(), if the solution is selected. - static const float kSlicedCopyResourceInflation = 1.8; - for (int i = 0; i < slice_start_times.size(); ++i) { - float original_copy_resource = copy_resource_per_slice[i]; - float new_copy_resource = original_copy_resource; - if (slice_start_times.size() > 1) { - // This is a hack that makes us more conservative about using sliced - // prefetching vs unsliced prefetching. - new_copy_resource = original_copy_resource * kSlicedCopyResourceInflation; - VLOG(5) - << "Inflating required copy resources DoWeHaveEnoughCopyResource() " - "slice check from " - << original_copy_resource << " to " << new_copy_resource; - } - specs.push_back( - {slice_start_times[i], prefetch_end_time, new_copy_resource}); - } - - auto specs_to_string = [&specs]() { - return absl::StrCat( - "[ ", - absl::StrJoin(specs, ", ", - [](std::string* out, - const AsynchronousCopyResource::ResourceSpec& spec) { - absl::StrAppend(out, "{exclusive start: ", - spec.exclusive_start_time, - ", end: ", spec.end_time, - ", resource: ", spec.resource, "}"); - }), - " ]"); - }; - - VLOG(5) << "Checking for enough copy resources for: " << specs_to_string(); - if (!async_copy_resource.HasEnoughResourceMultiCheck(specs)) { - VLOG(4) << "Not enough copy resources for " << specs_to_string(); - return false; - } - return true; -} - -// We compute a map from indices in chunk_candidates to indices in a -// SliceProposalCollection. Since the indices of chunk_candidates correspond to -// slice start times order, and SliceProposalCollections are always sorted in -// offset order, the mapping allows us to get the sizing details of a slice at a -// specific slice time. -absl::flat_hash_map GetCandidateToProposalIndexMap( - const std::vector& chunk_candidates) { - std::vector> sorted_offset_candidate_index_pairs; - sorted_offset_candidate_index_pairs.reserve(chunk_candidates.size()); - for (int64_t chunk_candidate_index = 0; - chunk_candidate_index < chunk_candidates.size(); - ++chunk_candidate_index) { - sorted_offset_candidate_index_pairs.push_back(std::make_pair( - chunk_candidates[chunk_candidate_index].offset, chunk_candidate_index)); - } - absl::c_sort(sorted_offset_candidate_index_pairs); - - absl::flat_hash_map candidate_to_proposal_index_map; - for (int64_t offset_index = 0; - offset_index < sorted_offset_candidate_index_pairs.size(); - ++offset_index) { - int64_t chunk_candidate_index = - sorted_offset_candidate_index_pairs[offset_index].second; - candidate_to_proposal_index_map[chunk_candidate_index] = offset_index; - } - - return candidate_to_proposal_index_map; -} - -} // namespace - -AllocationResult MsaAlgorithm::CheckPrefetchFit(bool for_sliced_solution, - PrefetchContext& context) { - SlicedBufferInterval* sliced_buffer_interval = - context.GetMutableWorkingIntervals(for_sliced_solution).sliced.get(); - - if (for_sliced_solution) { - CHECK(context.slice_proposal_collection); - CHECK_EQ(context.slice_proposal_collection->size(), - sliced_buffer_interval->num_slices()); - } - - // Update the prefetch start time in our working solution. - std::vector exclusive_slice_start_times = - SlicedPrefetchStartTimePicker::Pick( - sliced_buffer_interval->num_slices(), - context.exclusive_prefetch_start_time, context.prefetch_end_time, - [&](int64_t exclusive_start_time, - int64_t exclusive_end_time) -> float { - return options_.prefetch_interval_picker->GetLogicalIntervalElapsed( - exclusive_start_time, exclusive_end_time); - }, - [&](int64_t lhs_time, int64_t rhs_time) -> bool { - return hlo_live_range_.flattened_instruction_sequence() - .instructions()[lhs_time] - ->parent() == - hlo_live_range_.flattened_instruction_sequence() - .instructions()[rhs_time] - ->parent(); - }); - CHECK_EQ(sliced_buffer_interval->num_slices(), - exclusive_slice_start_times.size()); - sliced_buffer_interval->UpdateExclusiveSliceStartTimes( - exclusive_slice_start_times); - VLOG(4) << AlternateMemoryAllocationAttemptToString(for_sliced_solution, - context); - - // Check if all slices have the same start time. If so, we might as well - // resort to a full copy. - if (for_sliced_solution && - absl::c_all_of( - exclusive_slice_start_times, [&](int64_t slice_start_time) { - return slice_start_time == exclusive_slice_start_times.front(); - })) { - return AllocationResult::kAllSlicesHaveTheSameStartTime; - } - - // Check that we have enough copy resource for the prefetching. - std::vector copy_resource_per_slice_sorted_by_start_time; - // If there is a preferred prefetch time due to a loop optimized allocation, - // we already keep track of the prefetch resources there, so skip tracking - // resources here. - if (context.request->preferred_prefetch_time) { - copy_resource_per_slice_sorted_by_start_time = - std::vector(exclusive_slice_start_times.size(), 0.0); - } else if (for_sliced_solution) { - // In a sliced setting, we don't yet know when each slice will be - // prefetched. Given the proposed slice times, the most conservative copy - // resource check we can make is to assume that larger slices are started - // at earlier times, i.e., they have more time to complete. That is the - // check we will make here. Once, we've decided when each slice will be - // prefetched, we can do an exact check below. - // - // We start by computing the amount of copy resources needed for each slice, - // if larger slices are started at earlier times. - copy_resource_per_slice_sorted_by_start_time = - GetCopyResourcesSortedDescending(options_, - *context.slice_proposal_collection); - } else { - copy_resource_per_slice_sorted_by_start_time.push_back( - CopyResourceForShape(options_, *context.full_shape)); - } - CHECK_EQ(sliced_buffer_interval->num_slices(), - copy_resource_per_slice_sorted_by_start_time.size()); - - // Check if the copies we would add for the prefetch would violate copy - // ordering. - if (options_.enforce_prefetch_fifo_order && - absl::c_any_of(exclusive_slice_start_times, - [&](int64_t slice_start_time) { - return async_copy_ordering_.ViolatesOrdering( - slice_start_time, context.prefetch_end_time); - })) { - VLOG(4) << "This would violate asynchronous copy ordering."; - return AllocationResult::kFailViolatesAsyncCopyResource; - } - - // Check if the copies we would add for the prefetch violate the maximum - // number of outstanding async copies. - for (int i = 0; i < exclusive_slice_start_times.size(); ++i) { - if (ViolatesMaximumOutstandingAsyncCopies( - exclusive_slice_start_times[i], context.prefetch_end_time, - /*is_prefetch=*/true, context.extra_async_copy_limit, i)) { - VLOG(4) << "This would violate the outstanding async copy limit."; - return AllocationResult::kFailOutOfAsyncCopies; - } - } - - if (!DoWeHaveEnoughCopyResource(exclusive_slice_start_times, - context.prefetch_end_time, - copy_resource_per_slice_sorted_by_start_time, - prefetch_async_copy_resource_)) { - return AllocationResult::kFailViolatesAsyncCopyResource; - } - - // Check if we can find a place in alternate memory for the prefetch. - std::vector chunk_candidates = FindBestChunkCandidates( - *context.request, context.request->preferred_offset, - sliced_buffer_interval); - CHECK(chunk_candidates.empty() || - chunk_candidates.size() == sliced_buffer_interval->num_slices()); - std::string prefetch_picker_debug_string; - if (VLOG_IS_ON(4)) { - prefetch_picker_debug_string = - options_.prefetch_interval_picker->ToDebugString(); - } - if (for_sliced_solution && !chunk_candidates.empty()) { - // We're trying a sliced solution. So, if FindBestChunkCandidates() found a - // solution, each slice should have its own chunk candidate. - CHECK_EQ(chunk_candidates.size(), sliced_buffer_interval->num_slices()); - // We need a mapping from chunks in chunk_candidates to slice proposals in - // context.slice_proposal_context. - absl::flat_hash_map candidate_to_proposal_index_map = - GetCandidateToProposalIndexMap(chunk_candidates); - - // Create slice decisions, sorted by time. - std::vector slice_decisions_sorted_by_start_time; - for (int64_t slice_time = 0; - slice_time < sliced_buffer_interval->num_slices(); ++slice_time) { - const SliceProposal& proposal = context.slice_proposal_collection->at( - candidate_to_proposal_index_map[slice_time]); - copy_resource_per_slice_sorted_by_start_time[slice_time] = - CopyResourceForShape(options_, proposal.slice_shape); - slice_decisions_sorted_by_start_time.push_back(SliceDecision{ - chunk_candidates[slice_time], exclusive_slice_start_times[slice_time], - proposal, copy_resource_per_slice_sorted_by_start_time[slice_time]}); - } - - // Check that we have enough copy resources for all the slice decisions. - if (!DoWeHaveEnoughCopyResource( - exclusive_slice_start_times, context.prefetch_end_time, - copy_resource_per_slice_sorted_by_start_time, - prefetch_async_copy_resource_)) { - return AllocationResult::kFailViolatesAsyncCopyResource; - } - - // Construct MsaBufferInterval-Chunk pairs that are appropriate for pending - // chunks, as described in PrefetchContext::SlicedSolution. - std::vector> slices_for_pending_chunks; - slices_for_pending_chunks.reserve(sliced_buffer_interval->num_slices()); - Chunk final_chunk = Chunk::FromOffsetSize( - absl::c_min_element( - chunk_candidates, - [](const Chunk& a, const Chunk& b) { return a.offset < b.offset; }) - ->offset, - absl::c_accumulate( - chunk_candidates, 0, - [](int64_t sum, const Chunk& chunk) { return sum + chunk.size; })); - MsaBufferInterval final_buffer_interval{ - context.request->allocation_value->value(), - /*size=*/final_chunk.size, - /*start=*/ - ExclusiveToInclusiveStartTime(exclusive_slice_start_times.back()), - /*end=*/context.request->end_time, - /*colocations=*/ - sliced_buffer_interval->full_buffer_interval().colocations, - /*need_allocation=*/true}; - for (int64_t slice_time = 0; - slice_time < sliced_buffer_interval->num_slices(); ++slice_time) { - const Chunk& chunk = chunk_candidates[slice_time]; - int64_t inclusive_start_time = ExclusiveToInclusiveStartTime( - exclusive_slice_start_times[slice_time]); - if (inclusive_start_time == - ExclusiveToInclusiveStartTime(exclusive_slice_start_times.back())) { - // This and the following chunks will be merged into the final chunk. - // Note, it's possible for more than one slice to start at the same - // time. - break; - } - CHECK_LT(inclusive_start_time, ExclusiveToInclusiveStartTime( - exclusive_slice_start_times.back())); - slices_for_pending_chunks.push_back(std::make_pair( - MsaBufferInterval{ - context.request->allocation_value->value(), - /*size=*/chunk.size, - /*start=*/inclusive_start_time, - /*end=*/exclusive_slice_start_times.back(), - // We only use the final_buffer_interval for colocations because - // slices start at different offsets, and the colocation - // infrastructure expects all colocated buffers to start at the - // same offset. - /*colocations=*/{}, - /*need_allocation=*/true, - }, - chunk)); - } - slices_for_pending_chunks.push_back( - std::make_pair(final_buffer_interval, final_chunk)); - - context.sliced_solution = PrefetchContext::SlicedSolution{ - std::move(slice_decisions_sorted_by_start_time), - std::move(slices_for_pending_chunks), - prefetch_picker_debug_string, - }; - return AllocationResult::kSuccess; - } else if (!chunk_candidates.empty()) { - // We're trying an unsliced solution. So, if FindBestChunkCandidates() found - // a solution, there must be only 1 chunk for it. - CHECK_EQ(chunk_candidates.size(), 1); - CHECK_EQ(copy_resource_per_slice_sorted_by_start_time.size(), 1); - context.unsliced_solution = PrefetchContext::UnslicedSolution{ - chunk_candidates.front(), - copy_resource_per_slice_sorted_by_start_time.front(), - prefetch_picker_debug_string, - }; - return AllocationResult::kSuccess; - } - - // Mark the out of memory start with the prefetch start time so that we don't - // explore prefetch start times earlier than this point. If a sliced prefetch - // doesn't fit at a given time, an unsliced prefetch will not fit either. - // Thus, if we are considering a sliced prefetch for the current request, - // we can only update out_of_mem_start when we check with slices. - if (for_sliced_solution || !context.slice_proposal_collection) { - CHECK_GT(exclusive_slice_start_times.size(), 0); - context.exclusive_out_of_mem_start = std::max( - context.exclusive_out_of_mem_start ? *context.exclusive_out_of_mem_start - : -1, - exclusive_slice_start_times.front()); - } - - VLOG(4) << "Out of memory."; - return AllocationResult::kFailOutOfMemory; -} - -std::string MsaAlgorithm::AlternateMemoryAllocationAttemptToString( - bool for_sliced_solution, const PrefetchContext& context) const { - const SlicedBufferInterval* sliced_buffer_interval = - context.GetWorkingIntervals(for_sliced_solution).sliced.get(); - - std::vector slice_times; - std::vector estimated_slice_prefetch_end_times; - - for (int i = 0; i < sliced_buffer_interval->num_slices(); ++i) { - slice_times.push_back(absl::StrCat( - "[", sliced_buffer_interval->IntervalForMakeFreeChunks(i).start, ", ", - sliced_buffer_interval->full_buffer_interval().end, ")")); - if (context.slice_proposal_collection) { - estimated_slice_prefetch_end_times.push_back( - options_.prefetch_interval_picker->EstimatedPrefetchEndTime( - context.slice_proposal_collection->at(i).slice_shape, - sliced_buffer_interval->IntervalForMakeFreeChunks(i).start, - context.prefetch_end_time)); - } else { - estimated_slice_prefetch_end_times.push_back( - options_.prefetch_interval_picker->EstimatedPrefetchEndTime( - *context.full_shape, - sliced_buffer_interval->IntervalForMakeFreeChunks(i).start, - context.prefetch_end_time)); - } - } - - return absl::StrCat( - "Trying alternate memory allocation. Slice times = { ", - absl::StrJoin(slice_times, ", "), " }. Estimated prefetch end times = { ", - absl::StrJoin(estimated_slice_prefetch_end_times, ", "), " }"); -} - -std::optional MsaAlgorithm::FindBestChunkCandidate( - const AllocationRequest& request, const AliasedOffset* preferred_offset, - MsaBufferInterval* alternate_mem_interval) const { - SlicedBufferInterval sliced_buffer_interval = - SlicedBufferInterval::CreateMutableInterval(*alternate_mem_interval); - std::vector chunks = FindBestChunkCandidates(request, preferred_offset, - &sliced_buffer_interval); - CHECK_LE(chunks.size(), 1); - if (chunks.empty() || chunks[0].chunk_end() > options_.max_size_in_bytes) { - return std::nullopt; - } - return chunks[0]; -} - -std::vector MsaAlgorithm::FindBestChunkCandidates( - const AllocationRequest& request, const AliasedOffset* preferred_offset, - SlicedBufferInterval* alternate_mem_interval) const { - int64_t end_time = request.end_time; - if (!preferred_offset) { - // First find the earliest use that is the same or later than the end time. - const auto& use_times = request.all_use_times; - auto use_time_it = absl::c_lower_bound(use_times, end_time); - CHECK(use_time_it != use_times.end()); - int64_t earliest_use = *use_time_it; - auto earliest_use_it = use_time_it; - - // Then find the latest use that can be allocated contiguously without - // copies. - const Shape& shape = request.allocation_value_to_update->position().shape(); - const int64_t shape_size = GetShapeSizeBytes(options_.cost_analysis, shape); - for (; - (use_time_it + 1) != use_times.end() && - options_.prefetch_interval_picker->CanAllocateInAlternateMemoryNoCopy( - shape_size, *use_time_it, *(use_time_it + 1)); - ++use_time_it) { - } - CHECK(use_time_it != use_times.end()); - int64_t latest_contiguous_use_time = *use_time_it; - - // Find chunks that are as long living as possible. - std::vector last_chunk_candidates; - int64_t latest_matching_use = std::numeric_limits::min(); - (void)std::lower_bound( - earliest_use_it, std::next(use_time_it), -1, [&](int64_t use, int64_t) { - alternate_mem_interval->UpdateEndTime(use); - std::vector chunk_candidates = - FindChunkCandidates(*alternate_mem_interval); - if (chunk_candidates.empty()) { - return false; - } - int64_t max_chunk_end = - absl::c_max_element(chunk_candidates, [](const Chunk& c1, - const Chunk& c2) { - return c1.chunk_end() < c2.chunk_end(); - })->chunk_end(); - if (max_chunk_end <= options_.max_size_in_bytes) { - if (use > latest_matching_use) { - last_chunk_candidates = std::move(chunk_candidates); - latest_matching_use = use; - } - return true; - } - return false; - }); - if (!last_chunk_candidates.empty()) { - VLOG(3) << "FindBestChunkCandidates earliest use = " << earliest_use - << ", latest contiguous use = " << latest_contiguous_use_time - << ", use with available mem = " << latest_matching_use - << ", offsets = { " - << absl::StrJoin(last_chunk_candidates, ", ", - [](std::string* out, const Chunk& c) { - absl::StrAppend(out, c.offset); - }) - << " }"; - } - alternate_mem_interval->UpdateEndTime(end_time); - return last_chunk_candidates; - } - // If a preferred offset is given, try to find an allocation at that offset - // only. - alternate_mem_interval->UpdateEndTime(end_time); - std::vector chunk_candidates = - FindChunkCandidates(*alternate_mem_interval, preferred_offset->offset); - // Ensure that chunk candidates exist before querying min/max elements to - // prevent undefined behavior or segmentation faults when chunk_candidates is - // empty. - if (chunk_candidates.empty()) { - return {}; - } - int64_t candidates_start = - absl::c_min_element(chunk_candidates, [](const Chunk& c1, - const Chunk& c2) { - return c1.offset < c2.offset; - })->offset; - int64_t max_chunk_end = - absl::c_max_element(chunk_candidates, [](const Chunk& c1, - const Chunk& c2) { - return c1.chunk_end() < c2.chunk_end(); - })->chunk_end(); - if (candidates_start == preferred_offset->offset && - max_chunk_end <= options_.max_size_in_bytes) { - return chunk_candidates; - } - return {}; -} - -bool MsaAlgorithm::IsPositionColoredInAlternateMemory( - const HloPosition& position) const { - int64_t instruction_time = - hlo_live_range_.instruction_schedule().at(position.instruction); - return IsPositionColoredInAlternateMemoryAtTime(position, instruction_time); -} - -bool MsaAlgorithm::IsUseColoredInAlternateMemory(const HloUse& use) const { - HloPosition position = HloPosition{ - use.instruction->mutable_operand(use.operand_number), use.operand_index}; - HloValue& value = alias_analysis_.dataflow_analysis().GetUniqueValueAt( - position.instruction, position.index); - HloPosition defining_position = value.defining_position(); - int64_t use_time = hlo_live_range_.instruction_schedule().at(use.instruction); - return IsPositionColoredInAlternateMemoryAtTime(defining_position, use_time); -} - -bool MsaAlgorithm::IsPositionColoredInDefaultMemory( - const HloPosition& position) const { - int64_t instruction_time = - hlo_live_range_.instruction_schedule().at(position.instruction); - return IsPositionColoredInDefaultMemoryAtTime(position, instruction_time); -} - -bool MsaAlgorithm::IsUseColoredInDefaultMemory(const HloUse& use) const { - HloPosition position = HloPosition{ - use.instruction->mutable_operand(use.operand_number), use.operand_index}; - HloValue& value = alias_analysis_.dataflow_analysis().GetUniqueValueAt( - position.instruction, position.index); - HloPosition defining_position = value.defining_position(); - int64_t use_time = hlo_live_range_.instruction_schedule().at(use.instruction); - return IsPositionColoredInDefaultMemoryAtTime(defining_position, use_time); -} - -bool MsaAlgorithm::IsPositionColoredInAlternateMemoryAtTime( - const HloPosition& position, int64_t time) const { - // Check if the defining position of the corresponding value is colored in - // alternate memory at the time of the use or the defining position of an - // aliasing value is colored in alternate memory at the time of the use. - for (const HloValue* value : - alias_analysis_.GetUniqueBufferAt(position.instruction, position.index) - .values()) { - HloPosition defining_position = value->defining_position(); - const HloBuffer& buffer = alias_analysis_.GetUniqueBufferAt( - defining_position.instruction, defining_position.index); - auto reserved_allocations_it = - reserved_allocations_for_alt_mem_colorings_.find(&buffer); - if (reserved_allocations_it != - reserved_allocations_for_alt_mem_colorings_.end()) { - const auto& allocations = reserved_allocations_it->second; - auto it = - std::lower_bound(allocations.begin(), allocations.end(), time, - [](const std::unique_ptr& a, - int64_t t) { return a->end_time() < t; }); - if (it != allocations.end() && (*it)->start_time() <= time) { - return true; - } - } - } - return false; -} - -bool MsaAlgorithm::IsPositionColoredInDefaultMemoryAtTime( - const HloPosition& position, int64_t time) const { - // Check if the defining position of the corresponding value is colored in - // default memory at the time of the use or the defining position of an - // aliasing value is colored in default memory at the time of the use. - for (const HloValue* value : - alias_analysis_.GetUniqueBufferAt(position.instruction, position.index) - .values()) { - HloPosition defining_position = value->defining_position(); - const HloBuffer& buffer = alias_analysis_.GetUniqueBufferAt( - defining_position.instruction, defining_position.index); - auto default_memory_colorings_it = - default_memory_coloring_requirements_.find(&buffer); - if (default_memory_colorings_it != - default_memory_coloring_requirements_.end()) { - const auto& intervals = default_memory_colorings_it->second; - auto it = std::lower_bound(intervals.begin(), intervals.end(), time, - [](const TimeInterval& interval, int64_t t) { - return interval.inclusive_end_time < t; - }); - if (it != intervals.end() && it->inclusive_start_time <= time) { - return true; - } - } - } - return false; -} - -int64_t AsynchronousCopyResource::GetScaledIntegerResource( - float resource) const { - float scaled_value = resource * kCopyResourceIntScale; - if (scaled_value > std::numeric_limits::max()) { - LOG(WARNING) << "Scaled value " << scaled_value - << " is greater than the maximum int64_t value " - << std::numeric_limits::max(); - return std::numeric_limits::max(); - } - auto scaled_value_int = static_cast(scaled_value); - return scaled_value_int; -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.h b/third_party/xla/xla/service/memory_space_assignment/algorithm.h deleted file mode 100644 index 9a45cfe3a6e1a3..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.h +++ /dev/null @@ -1,1956 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALGORITHM_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALGORITHM_H_ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// TODO(b/210891274): Use btree_map after build issue in Windows is resolved. -#if defined(__GNUC__) || defined(__clang__) -#include "absl/container/btree_map.h" -#endif -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/container/linked_hash_map.h" -#include "absl/hash/hash.h" -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/call_graph.h" -#include "xla/service/decision.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo.pb.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/allocation_value.h" -#include "xla/service/memory_space_assignment/live_range_util.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/service/memory_space_assignment/utils.h" -#include "xla/shape.h" -#include "xla/shape_tree.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -// A struct representing use intervals. -struct UseInterval { - int64_t first_use_time; - int64_t last_use_time; -}; - -struct TimeInterval { - int64_t inclusive_start_time; - int64_t inclusive_end_time; - - bool operator<(const TimeInterval& other) const; - - std::string ToString() const; -}; - -// A struct representing an asynchronous copy with its logical start and end -// time (time that copy done is scheduled), the resource this copy would use, -// its destination memory space, and a unique ID. -struct AsynchronousCopy { - int64_t exclusive_start_time; - int64_t end_time; - float resource; - MemorySpace destination; - int64_t id; - - std::tuple AsTuple() const { - return std::make_tuple(exclusive_start_time, end_time, resource, - destination, id); - } -}; - -// Represents a context for allocating a segment of an AllocationValue. -// AllocationValue typically provides enough information to allocate the entire -// live range of the AllocationValue, since all segments update only the -// AllocationSequence belonging to the AllocationValue. However, in cases of -// synchronous memory op conversion (e.g., copy, slice, etc.), we also need -// to modify the AllocationSequence of the AllocationValue produced at the -// synchronous memory op's output. This struct provides a context for allocating -// a segment of an AllocationValue, specifying the uses of the AllocationValue -// that we are processing, the index of the use that we are processing in the -// AllocationValue's uses vector, the index of the AllocationValue in -// Span, whose allocation sequence we will update, and whether -// the use is only processed to extend the lifetime of its operand's allocation, -// and the use will not receive a new allocation. -struct AllocationSegmentContext { - // The uses of the AllocationValue that we are processing. - const std::vector* uses; - // The index of the use that we are processing in the AllocationValue's - // AllocationValue::uses vector. - int use_idx; - // Index of the AllocationValue in allocation_values that is being processed - // in AllocateAllocationValues(), whose allocation sequence we will be - // updated. - int allocation_value_to_update_idx; - // If true, the use is only processed to extend the lifetime of its operand's - // allocation, and the use will not receive a new allocation. - bool only_extend_existing_allocation; -}; - -// Compare asynchronous copies such that an earlier start time has the same or -// earlier end time and an earlier end time has the same or earlier start time. -bool operator<(const AsynchronousCopy& a, const AsynchronousCopy& b); - -bool operator==(const AsynchronousCopy& a, const AsynchronousCopy& b); -bool operator!=(const AsynchronousCopy& a, const AsynchronousCopy& b); - -// Helper class to enforce asynchronous copy ordering. If the appropriate option -// is enabled, we only allow asynchronous copies that are pipelined: if an -// asynchronous copy ends earlier than another asynchronous copy, it must start -// the same time or earlier than the other asynchronous copy; and if an -// asynchronous copy starts earlier than another asynchronous copy, it must end -// the same time or earlier than the other asynchronous copy. -class AsynchronousCopyOrdering { - public: - AsynchronousCopyOrdering() = default; - - // Adds an asynchronous copy. - void AddCopy(const AsynchronousCopy& copy); - - // Removes an asynchronous copy. CHECKs that it is removed. - void RemoveCopy(const AsynchronousCopy& copy); - - // Returns true if the addition of an asynchronous copy in the given time - // interval would violate the asynchronous copy ordering. E.g., consider the - // following scenario: - // CS CD - // already committed async copy: +-----------+ - // new async copy: +--------+ - // - // The new asynchronous copy would violate the ordering guarantee because the - // copy start is after an already committed asynchronous copy while its copy - // done is before the committed copy. - bool ViolatesOrdering(int64_t exclusive_start_time, int64_t end_time) const; - - private: - // We use this data structure for keys into the map that has a custom - // comparator for the ordering guarantees. - struct Interval { - int64_t exclusive_start_time; - int64_t end_time; - - // We allow multiple prefetches that have one or both of the same start and - // end times. std::map considers two values as equal if neither are less - // than the other. Using this comparator, we can ensure that the only - // intervals that evaluate to be equal are those with the same start and end - // times or those with intervals that violate the FIFO order. - bool operator<(const Interval& other) const { - return (exclusive_start_time < other.exclusive_start_time && - end_time <= other.end_time) || - (exclusive_start_time <= other.exclusive_start_time && - end_time < other.end_time); - } - }; - // Stores asynchronous copies in a tree set respecting the pipelining order. - std::map> ranges_; -}; - -// Helper class to enforce asynchronous copy resources by keeping track of -// available copy bandwidth and elapsed times of overlapped operations. It -// maintains a list of initial resources that correspond to the elapsed times of -// overlapped operations. As asynchronous copies are added, the available -// resource is subtracted to keep track of the current state. -class AsynchronousCopyResource { - public: - // A specification of needed asynchronous copy resources. - struct ResourceSpec { - int64_t exclusive_start_time; - int64_t end_time; - float resource; - }; - - AsynchronousCopyResource() = default; - - // The constructor needs the initial resources. - explicit AsynchronousCopyResource(absl::Span initial_resources) - : initial_resources_(initial_resources.begin(), initial_resources.end()), - delay_(initial_resources.size(), 0) { - for (int i = 0; i < initial_resources.size(); ++i) { - initial_resources_scaled_.push_back( - GetScaledIntegerResource(initial_resources[i])); - } - } - - // Adds the given asynchronous copy and updates the current resources. CHECK - // fails if there aren't enough resources to satisfy this copy (the caller - // should use HasEnoughResource first to ensure there is enough resource). - void AddCopy(const AsynchronousCopy& copy); - - // Removes the given copy and frees the resource. - void RemoveCopy(const AsynchronousCopy& copy); - - // Returns true if a copy with the given start and end times and resource can - // be satisfied. - bool HasEnoughResource(int64_t exclusive_start_time, int64_t end_time, - float resource); - - // Returns true if a set of copy specifications can be satisfied in the - // order specified. - bool HasEnoughResourceMultiCheck(const std::vector& specs); - - int64_t GetScaledIntegerResource(float resource) const; - - float GetDescaledFloatResource(int64_t scaled_resource) const { - return scaled_resource / kCopyResourceIntScale; - } - - // This is only used for debugging and testing purposes, it returns the - // currently available resource at each logical time. - std::vector GetCurrentResources() const { - std::vector current_resources(initial_resources_.begin(), - initial_resources_.end()); - for (int i = 0; i < current_resources.size(); ++i) { - current_resources[i] -= - std::min(current_resources[i], GetDescaledFloatResource(delay_[i])); - } - return current_resources; - } - - // A useful debugging tool for printing several pieces of information about - // AsynchronousCopyResource. - std::string Dump(int64_t start_time, int64_t end_time, - MemorySpace memory_space_filter) const; - - // The scale factor to convert a float resource to an integer resource. Note - // that is a power of 2 to avoid introducing noise when casting the scaled - // value to an int64_t. - static constexpr int64_t kCopyResourceIntScale = 1ULL << 40; - - private: - // Internal helper method to implement adding/removing/checking resources. - // ConsumeResource() may modify delay_. If delay_changes is not null, - // for any change to delay_[i], {i, delay_[i]} will be added to - // delay_changes, allowing callers to undo any modifications by iterating over - // the vector in reverse order. - bool ConsumeResource( - int64_t exclusive_start_time, int64_t end_time, int64_t resource, - std::vector>* delay_changes = nullptr, - int64_t resource_to_free = 0.0); - - // Same as the public RemoveCopy except it works on the async_copies_ - // iterator. Assumes copy_it points to the last copy for its start time; - // otherwise the public RemoveCopy method is supposed to temporarily remove - // these later copies that share the same start time before removing the - // requested copy. - void RemoveCopy(std::list::iterator& copy_it); - - // We maintain a linked list of asynchronous copies sorted by the start times. - // This allows us to efficiently find the copy that starts right after another - // one because adding a copy might push a copy further into the future. - std::list async_copies_; -// To make the lookups into async_copies_ more efficient, we also maintain a -// binary tree that is indexed by the start time, containing iterators into -// async_copies_. -// TODO(b/210891274): Use btree_map after build issue in Windows is resolved. -#if defined(__GNUC__) || defined(__clang__) - absl::btree_map::iterator> - async_copy_time_map_; -#else - std::map::iterator> async_copy_time_map_; -#endif - std::vector initial_resources_; - std::vector initial_resources_scaled_; - std::vector delay_; - // A vector of pairs of (time, delay) used by - // HasEnoughResourceMultiCheck(), stored here to avoid reallocations. - std::vector> delay_changes_; -}; - -// Helper class to compute a minimal fingerprint of an HloInstruction and it's -// operand shapes for MSA. -// Use a cache to avoid recomputing the fingerprint of the shape for the same -// instruction more than once. -template -class MsaInstructionFingerprint { - public: - explicit MsaInstructionFingerprint( - const HloInstruction* instruction, - absl::flat_hash_map* - instruction_shape_hash_cache = nullptr) - : inst_(instruction), - instruction_shape_hash_cache_(instruction_shape_hash_cache) {} - - template - friend H AbslHashValue(H h, const MsaInstructionFingerprint& fp) { - auto get_or_compute_shape_hash = [&fp](const HloInstruction* inst) { - if (fp.instruction_shape_hash_cache_ != nullptr) { - auto it = fp.instruction_shape_hash_cache_->find(inst); - if (it != fp.instruction_shape_hash_cache_->end()) { - return it->second; - } - } - HashType shape_hash = absl::HashOf(inst->shape()); - if (fp.instruction_shape_hash_cache_ != nullptr) { - fp.instruction_shape_hash_cache_->insert({inst, shape_hash}); - } - return shape_hash; - }; - for (const HloInstruction* operand : fp.inst_->operands()) { - h = H::combine(std::move(h), get_or_compute_shape_hash(operand)); - } - return H::combine(std::move(h), fp.inst_->opcode(), - fp.inst_->operand_count(), - get_or_compute_shape_hash(fp.inst_)); - } - - private: - const HloInstruction* inst_; - absl::flat_hash_map* - instruction_shape_hash_cache_; -}; - -// This class inherits from GlobalDecreasingSizeBestFitHeap with a notion of -// maximum size. -// -// Note: Memory space assignment (MSA) creates an MsaAlgorithm object and passes -// it to the HeapSimulator. The HeapSimulator calls Alloc(), Free() and -// ShareWith() on the MsaAlgorithm object to create buffer intervals (populate -// buffer_intervals_), these methods are inherited from -// GlobalDecreasingSizeBestFitHeap. The HeapSimulator finally calls the Finish() -// method which is overridden in this class. -class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap { - public: - MsaAlgorithm(HloModule* module, AllocationSequence* allocations, - const Options& options, const HloAliasAnalysis& alias_analysis, - const AliasInfo* alias_info, const HloLiveRange& hlo_live_range); - - // Allocates a buffer in preferred memory with whole program lifetime and - // enables prefetching prefetch_candidate from default memory across program - // boundaries. - void AllocateCrossProgramPrefetchBuffer( - HloModule* module, const MsaBufferInterval& prefetch_candidate); - - // Given an HloValue, returns a group of HloValues that need to be processed - // jointly. Normally, HloValues can be processed individually. However, in - // case we are trying to replace synchronous copies, we need to jointly - // process all values that are produced or consumed by a synchronous memory - // call instruction. - std::vector GenerateJointProcessedValues( - const HloValue* entrance_value); - - // Updates sorted_sync_copy_replacement_candidates_ with synchronous copy - // instructions that connect the given joint processed values, and meet the - // conditions in IsReplaceableSyncCopyCandidate(). - void UpdateSyncDataMovementCandidatesForJointProcessedValues( - const std::vector& joint_processed_values); - - // Returns true if repack_allocation_blocks_ includes an AllocationBlock - // belonging to a converted synchronous memory operations. - bool RepackAllocationsIncludeConvertedSyncMemOp(); - - absl::StatusOr> Finish() override; - - // Testing wrapper for FindAliases. - static void FindAliasesForTesting( - std::vector* allocation_values, - bool has_async_pipelined_while_loops = false); - - // Block prefetching is an MSA feature that allows processing all prefetches - // in one pass within a block of memory space in the alternate memory. This - // guarantees FIFO ordering of all prefetches and allows for more aggressive - // prefetching i.e. allowing for bandwidth saturation. - // - // 1) Prefetches are copy-like operations and generate a new HloValue. - // 2) For compiler-inserted block prefetches: - // A) The prefetch done, and everything that aliases with the prefetch - // source and comes after the prefetch done will now alias with the new - // HloValue and get alternate memory. - // B) Everything that aliases with the prefetch source and comes before the - // prefetch done will get default memory. - // 3) For user-inserted block prefetches: - // A) The prefetch done and everything that aliases with it will get - // alternate memory. - // B) Everything that aliases with the source of the prefetch will get - // default memory. - - // Processes existing, explicit block prefetched copy start/done instructions. - // Such instructions are inserted before MSA. MSA just needs to schedule them. - // - // REQUIRED: Scoped vmem must be allocated at offset 0 at the time this method - // is called. - absl::Status AllocateAndScheduleExistingBlockPrefetches( - int64_t block_prefetching_starting_offset); - - // Create, allocate and schedule new block prefetches, by adding async copies - // or asyncifying DMA ops like slice, dynamic-slice. - // - // REQUIRED: Scoped vmem must be allocated at offset 0 at the time this method - // is called. - absl::Status CreateNewBlockPrefetches( - int64_t block_prefetching_starting_offset); - - // Creates colocated allocations for values aliased to existing block - // prefetches and finalizes them. - void ColocateAndFinalizeValuesAliasedToExistingBlockPrefetches( - const HloValue* prefetch_done_value, const HloBuffer& buffer, - const Chunk& chunk_candidate, int64_t buffer_size, - AllocationBlock* first_colocated_repack_allocation, - const absl::flat_hash_map& - instruction_schedule, - const absl::flat_hash_map& - value_to_use_intervals, - absl::flat_hash_map& - prefetch_done_value_to_prefetch_start_instruction, - std::vector& prefetch_end_times); - - // Returns the maximum amount of scoped memory that is reserved at any time in - // the program. - int64_t MaxScopedMemorySize(); - - // Finds and returns the earliest block prefetch start time subject to the - // following constraints: - // - We can reserve a chunk of size buffer_size. - // - The chunk's start_time lies between [definition_time, first_use_time]. - // - The chunk's end_time is equal to last_use_time. - // - The chunk's ending offset is less than or equal to - // block_prefetching_limit_bytes. - // - The number of concurrent prefetches is less than or equal to - // max_in_flight_prefetches_allowed. - // If no such prefetch_start_time exists, returns std::nullopt. - std::optional EarliestBlockPrefetchStartTime( - int64_t previous_block_start_time, int64_t definition_time, - int64_t first_use_time, int64_t last_use_time, int64_t buffer_size, - int64_t block_prefetching_limit_bytes, - int64_t max_in_flight_prefetches_allowed, - std::vector& copy_done_schedule_before_times, - std::vector& block_prefetch_allocation_end_times); - - // Calculates the latest schedule time of the non-constant source operands of - // an instruction, ensuring prefetching does not start before inputs are - // available. Returns std::nullopt if there are no non-constant source - // operands. - std::optional GetLatestSourceOperandScheduleTime( - HloInstruction* instruction, - const absl::flat_hash_map& - instruction_schedule) const; - - protected: - // Given a buffer interval, returns the colocated intervals. Unlike the - // similar GlobalDecreasingSizeBestFitHeap::GetTransitiveColocations, it - // returns the colocated intervals sorted by scheduled time. - std::vector GetSortedColocatedIntervals( - const MsaBufferInterval& interval) const; - - // Given a MsaBufferInterval, creates AllocationValue objects and - // corresponding AllocationSequences and appends them into - // allocation_sequence_list_. - void CreateAllocationValues( - const MsaBufferInterval& buffer_interval, - std::vector& allocation_values) const; - - // Given colocated intervals, populates allocation_values with the - // corresponding AllocationValue objects. - virtual void CreateAllocationValuesFromColocatedIntervals( - absl::Span colocated_intervals, - std::vector& allocation_values); - - AllocationSequence* allocations() { return allocations_; } - const Options& options() const { return options_; } - const HloAliasAnalysis& alias_analysis() { return alias_analysis_; } - const HloLiveRange& hlo_live_range() { return hlo_live_range_; } - - // Runs a feature that attempts to expand the size of scoped alternate memory - // allocations to the largest contiguous open space available. - void ExtendScopedAlternateMemoryAllocations(); - - // The same as FindBestChunkCandidate() but allocates the request in slices. - // The ith returned chunk should be allocated at slice time i. - std::vector FindBestChunkCandidates( - const AllocationRequest& request, const AliasedOffset* preferred_offset, - SlicedBufferInterval* alternate_mem_interval) const; - - // Virtual wrapper for FindChunkCandidates to allow testing subclasses to - // simulate candidate placement results. - virtual std::vector FindChunkCandidates( - const SlicedBufferInterval& sliced_buffer_interval, - int64_t preferred_offset) const { - return GlobalDecreasingSizeBestFitHeap::FindChunkCandidates( - sliced_buffer_interval, preferred_offset); - } - - // Convenience overload for FindChunkCandidates without a preferred offset. - std::vector FindChunkCandidates( - const SlicedBufferInterval& sliced_buffer_interval) const { - return FindChunkCandidates(sliced_buffer_interval, -1); - } - - private: - // Pins all scalar buffers in alternate memory. If a buffer has DMA like - // uses that can be asyncified, we need to make sure the buffer is live until - // the DMA done instruction. If the buffer is the source buffer of the DMA, - // MSA handles it. If the buffer is not the source buffer of the DMA, - // currently, it requires special handling. Scalars are the only buffers that - // have asyncifiable DMA uses in which they are not the source buffer of the - // DMA. In this case, we extend the live range of the scalar buffers to the - // end of the program as a temporary hack to make sure the buffer outlives any - // newly created async done instruction. - absl::Status PinScalarBuffersInAlternateMemory( - absl::flat_hash_set& scalar_buffers_to_pin_in_alt_mem); - - // Returns the set of scalar buffers that are pinned to the alternate - // memory. - absl::flat_hash_set GetScalarBuffersPinnedToAltMemory( - absl::Span sorted_buffer_intervals) const; - - // Helper method for PinScalarBuffersInAlternateMemory to pin a single scalar - // buffer's positions in the alternate memory. - absl::Status PinScalarBufferInAlternateMemory( - const HloBuffer* buffer, int64_t program_end_time, - const absl::flat_hash_map& - instruction_schedule); - - // We inherit AllocationBlock struct to attach the Allocation information to - // make importing repacked offsets easier. - struct RepackAllocationBlock : AllocationBlock { - Allocation* allocation; - }; - - // This struct contains mandatory memory assignments at a given time. E.g., an - // input's required memory assignment time would correspond to the definition - // time of the parameter instruction, and an output's time would correspond to - // the time of last use. - struct RequiredMemoryAssignment { - enum class Source { - kAliasedUse, - kUseNotAllowedInAlternateMemory, - kInefficientSite, - kLoopOptimizedParameterInWhileLoop, - kConditionalPhiOutput, - kConditionalComputationOutput, - kPositionNotAllowedInAlternateMemory, - kProgramOutput, - kConstantInstruction, - kProgramInput, - kBlockPrefetchSourceBuffer, - kBlockPrefetchSourceBufferUse, - }; - - MemorySpace memory_space; - int64_t time; - AliasedOffset* offset = nullptr; - // The source of the required memory assignment. Required memory assignments - // are created for various reasons, this field is used to identify the - // reason for the required memory assignment. It improves logging and helps - // debugging. - Source required_assignment_source; - - bool memory_space_and_offset_equal( - const RequiredMemoryAssignment& other) const { - return std::forward_as_tuple(memory_space, - offset ? offset->offset : -1) == - std::forward_as_tuple(other.memory_space, - other.offset ? other.offset->offset : -1); - } - - bool operator==(const RequiredMemoryAssignment& other) const { - return memory_space == other.memory_space && time == other.time && - offset == other.offset; - } - - bool operator!=(const RequiredMemoryAssignment& other) const { - return !(*this == other); - } - - std::string ToString() const; - static std::string SourceToString(Source source); - }; - - enum class BlockPrefetchCategory { - // Buffer must be pinned to default memory. - kPinToDefaultMemory, - // Buffer will be block prefetched only (no asyncification, e.g. program - // input). - kBlockPrefetchOnly, - // Buffer will be block prefetching and asyncified. - kAsyncConversionCandidate, - }; - - // Metadata associated with a block prefetching async conversion candidate. - struct AsyncConversionMetadata { - // The source HloBuffer of the async conversion candidate instruction. - const HloBuffer* source_buffer = nullptr; - // The non-trivial HloPosition corresponding to the source buffer operand. - HloPosition non_trivial_source_position; - }; - - // Encapsulates the results of analyzing all candidate positions for block - // prefetching. This analysis is a subcomponent of the BlockPrefetchContext. - // See BuildBlockPrefetchingContext() for more details. - struct BlockPrefetchCandidateAnalysis { - // Map from buffer to its block prefetch category. - absl::flat_hash_map - buffer_categories; - // Map from candidate buffer to its async conversion metadata. - absl::flat_hash_map - async_candidate_metadata; - // Map from buffer to its earliest defining HloValue. - absl::flat_hash_map - buffer_to_first_defining_value; - }; - - // This context contains initial block prefetching candidate analysis, and - // some other data structures used during the implementation of block - // prefetching. - struct BlockPrefetchContext { - // Initial analysis and categorization of candidate positions. - const BlockPrefetchCandidateAnalysis initial_analysis; - // List of candidate buffers sorted by their first use time and ID. - std::vector block_prefetched_buffers_list; - // Map from candidate buffer to its earliest use time in the schedule. - absl::flat_hash_map buffer_to_first_use_time; - // Set of instructions that are async conversion candidates. - absl::flat_hash_set - async_conv_candidate_instructions; - }; - - // Struct representing the result of analyzing uses for a block prefetch - // candidate. - struct BlockPrefetchCandidateUseAnalysis { - int64_t first_use_time = std::numeric_limits::max(); - bool all_uses_allowed_in_alternate_memory = true; - bool all_uses_are_synchronous = true; - const HloInstruction* first_use_instruction = nullptr; - }; - - // Encapsulates the block prefetch scheduling state that is maintained across - // iterations during CreateNewBlockPrefetches. - struct BlockPrefetchSchedulingState { - // Map from a buffer's defining source position (or async candidate - // position) to its pinned allocation in default memory, serving as the - // parent allocation for copies. - absl::flat_hash_map - buffer_defining_source_position_to_pinned_allocation; - // Sorted list of end times of currently active block prefetches, used to - // compute outstanding prefetches. - std::vector buffer_end_times; - // List of target times by which outstanding prefetches must complete. - std::vector prefetch_done_schedule_before_times; - // The start time of the previously scheduled prefetch, used to enforce - // FIFO ordering. - int64_t previous_start_time = -1; - }; - - // Builds the BlockPrefetchContext by running a 4-stage pipeline that - // categorizes block prefetching candidates and builds necessary data - // structures for executing on block prefetching. - // - // Stage 1: InitializeBlockPrefetchCandidateAnalysis() - // Identifies candidate buffers from options_.block_prefetched_positions, - // categorizes them (BlockPrefetchCategory), and extracts initial metadata. - // Stage 2: FilterAsyncConversionCandidates() - // Evaluates async conversion candidate buffers against source/destination - // DMA constraints to prevent alternate-memory to alternate-memory copies. - // Demotes failing candidates to kPinToDefaultMemory. May also impose - // kPinToDefaultMemory on sources to the async conversion candidates. - // Stage 3: FilterBlockPrefetchCandidatesByUseConstraints() - // Evaluates candidate buffers against alternate memory use compatibility - // and middle data movement positions. Demotes non-compliant candidates to - // kPinToDefaultMemory. - // Stage 4: BuildBlockPrefetchingContextHelper() - // Collects all non-demoted candidate buffers, computes their earliest use - // times, sorts them by (first_use_time, buffer_id) for deterministic - // scheduling order, and constructs the final BlockPrefetchContext. - const BlockPrefetchContext BuildBlockPrefetchingContext( - const absl::flat_hash_map& - instruction_schedule) const; - BlockPrefetchCandidateAnalysis InitializeBlockPrefetchCandidateAnalysis( - const absl::flat_hash_map& - instruction_schedule) const; - void FilterAsyncConversionCandidates( - BlockPrefetchCandidateAnalysis& initial_analysis) const; - void FilterBlockPrefetchCandidatesByUseConstraints( - BlockPrefetchCandidateAnalysis& initial_analysis, - const absl::flat_hash_map& - instruction_schedule) const; - BlockPrefetchContext BuildBlockPrefetchingContextHelper( - BlockPrefetchCandidateAnalysis initial_analysis, - const absl::flat_hash_map& - instruction_schedule) const; - - const HloBuffer* GetScheduledBuffer( - const HloPosition& position, - const absl::flat_hash_map& - instruction_schedule) const; - - BlockPrefetchCategory ClassifyBlockPrefetchBuffer( - const HloBuffer* buffer, const HloValue* first_defining_value) const; - - AsyncConversionMetadata ComputeBlockPrefetchingAsyncConversionMetadata( - HloInstruction* first_instruction, - const absl::flat_hash_map& - instruction_schedule) const; - - BlockPrefetchCandidateUseAnalysis ComputeBlockPrefetchCandidateUseAnalysis( - const HloBuffer* buffer, - const absl::flat_hash_map& - instruction_schedule, - const absl::flat_hash_set& - async_conv_candidate_instructions) const; - - bool HasDataMovementPositionsInMiddle( - const HloBuffer* buffer, const HloInstruction* first_instruction) const; - - // Attempts to allocate and schedule a block prefetch for a single buffer. - absl::StatusOr BlockPrefetchBuffer( - const HloBuffer* buffer, const HloValue* first_defining_value, - bool is_async_conversion_candidate, - std::optional async_non_trivial_source_position, - int64_t first_use_time, - const absl::flat_hash_map& - instruction_schedule, - int64_t block_prefetching_limit_bytes, - int64_t max_in_flight_prefetches_allowed, - const absl::flat_hash_set& - async_conv_candidate_instructions, - BlockPrefetchSchedulingState& scheduling_state); - - // Helper method for BlockPrefetchBuffer() to create or extend a pinned - // allocation in default memory for an HloPosition. - Allocation* CreatePinnedAllocationInDefaultMemory( - HloPosition position, int64_t first_use_time, - const absl::flat_hash_map& - instruction_schedule, - BlockPrefetchSchedulingState& scheduling_state, - AllocationSequence& allocations); - - // Helper method to create pinned allocations in alternate memory for - // HloPositions. Used by BlockPrefetchBuffer() and - // PinScalarBufferInAlternateMemory(). - void CreatePinnedAllocationsInAltMemoryForPositions( - absl::Span positions, const Chunk& chunk_candidate, - const absl::flat_hash_map& position_to_live_range, - const absl::flat_hash_map>& - position_to_uses, - const absl::flat_hash_set& - async_conv_candidate_instructions, - AllocationSequence& allocations, - std::vector& colocations, - // If set, overrides the end_time of the final position's pinned - // allocation. (Used to extend a scalar buffer's last allocation to the - // end of the program, as a temporary hack to make sure the buffer - // outlives any newly created async done instruction). If nullopt, the - // position's own live range end time is used. - std::optional last_position_end_time = std::nullopt); - - // Helper method to create mirrored allocations in alternate memory for - // HloPositions. Used by BlockPrefetchBuffer() and - // PinScalarBufferInAlternateMemory(). - absl::Status CreateMirroredAllocationsInAlternateMemory( - absl::Span positions, - const LiveRangeCalculator& calculator, - const absl::flat_hash_map& - instruction_schedule, - const absl::flat_hash_map>& - position_to_uses, - const absl::flat_hash_set& - async_conv_candidate_instructions, - AllocationSequence& allocations); - - // A struct that contains a pointer to loop-optimized allocation along with - // essential data about the loop itself. - struct LoopOptimizedAllocationInfo { - // The use_idx is the instruction index of the use within the loop. - int64_t use_index; - // The number of instructions in one iteration of the loop. We use use_index - // and loop_size to calculate when exactly to schedule a prefetch - // instruction. - int64_t loop_size; - // A pointer into an Allocation in loop_optimized_allocations_. - const Allocation* loop_optimized_allocation; - }; - - // A context object that is used to share state amongst the methods that - // implement Prefetch(). Prefetch tries to find both a sliced solution and an - // unsliced solution at the same time. We store both in this structure. - struct PrefetchContext { - // Prefetching is designed to operate on a SlicedBufferInterval that is - // backed by a standard MsaBufferInterval, even if the number of slices - // == 1. WorkingIntervals is used to store a SlicedBufferInterval and its - // backing MsaBufferInterval. - struct WorkingIntervals { - MsaBufferInterval full; - // sliced is a unique_ptr because it won't necessarily be initialized - // when the WorkingBufferIntervals are created, and there is no way to - // create an empty SlicedBufferInterval. - std::unique_ptr sliced; - }; - - struct SlicedSolution { - // When we talk about a slice, we think of spatial slices, where each - // slice is allocated at different times. The following example shows - // 3 slices that are used to form a contiguous buffer from [p0, p3] - // - // space - // ^ - // p3 | +-----------+ - // | | s2 | - // p2 | +---+-----------+ - // | | s1 | - // p1 | +-------+-------+ - // | | s0 | - // p0 | +-------+ - // +---|---|---|---|---|----> time - // t0 t1 t2 t3 t4 - std::vector slice_decisions_sorted_by_start_time; - - // In order to support colocated buffer calculations, we need to add a - // MsaBufferInterval-Chunk pair to pending_chunks_, such that: - // - The duration of the MsaBufferInterval is non-zero. - // - All slices have been allocated by the start of the MsaBufferInterval. - // - The MsaBufferInterval ends at the end time for all slices. - // - The Chunk covers the space allocated for all slices. - // - // In order to meet that requirement, - // we create MsaBufferInterval-Chunk pairs from - // slice_decisions_sorted_by_start_time that meet those requirement but do - // not cause any memory to be allocated in more than one Chunk at a time. - // The result is stored in slices_for_pending_chunks. - // - // The illustration below demonstrates how we would construct such - // MsaBufferInterval-Chunk pairs from the - // slice_decisions_sorted_by_start_time example above. - // - // space - // ^ - // p3 | +---+---+---+ - // | |c2 | | - // p2 | +---+---+ | - // | | c0 | c2 | - // p1 | +-------+ | - // | | | - // p0 | +-------+ - // +---|---|---|---|---|----> time - // t0 t1 t2 t3 t4 - std::vector> - slices_for_pending_chunks; - - // The prefetch_picker_debug_string will only be set with the appropriate - // VLOG level. - std::string prefetch_picker_debug_string; - }; - - struct UnslicedSolution { - Chunk chunk_candidate; // The chunk chosen for the solution. - float prefetch_resource; // The amount of required prefetch resource. - // The prefetch_picker_debug_string will only be set with the appropriate - // VLOG level. - std::string prefetch_picker_debug_string; - }; - - WorkingIntervals& GetMutableWorkingIntervals(bool for_sliced_solution) { - if (for_sliced_solution) { - return sliced_solution_intervals; - } - return unsliced_solution_intervals; - } - - const WorkingIntervals& GetWorkingIntervals( - bool for_sliced_solution) const { - if (for_sliced_solution) { - return sliced_solution_intervals; - } - return unsliced_solution_intervals; - } - - // Parameters to Prefetch(). - const AllocationRequest* request; - Allocation* prev_allocation_in_default_mem; - - // Intermediate calculations common to both the sliced and unsliced - // solutions. - int64_t exclusive_prefetch_start_time = -1; - int64_t prefetch_end_time = -1; - const Shape* full_shape; - int64_t extra_async_copy_limit = 0; - // As a compilation time optimization, store the prefetch start time where - // we have first seen out of memory. There is no point of exploring prefetch - // start times earlier than this point. - std::optional exclusive_out_of_mem_start = std::nullopt; - - // Data structures used to compute and store the sliced solution. - std::optional slice_proposal_collection = - std::nullopt; - WorkingIntervals sliced_solution_intervals; - std::optional sliced_solution; - - // Data structures used to compute and store the unsliced solution. - WorkingIntervals unsliced_solution_intervals; - std::optional unsliced_solution; - - // Indicates whether the prefetch is for a windowed prefetch. A window - // prefetch only prefetches a window worth of data. Its prefetch does not - // use sliced prefetch. - bool window_prefetch = false; - }; - - // Return true if the result belongs to a failure. - static bool result_is(AllocationResult result, AllocationResult failure) { - return static_cast(result) & static_cast(failure); - } - - // Mark (bitwise OR) a failure to the result. - static AllocationResult result_mark(AllocationResult failure, - AllocationResult& result) { - result = static_cast(static_cast(result) | - static_cast(failure)); - return result; - } - - // Return a string representation of a result that has at most a single - // failure. Consider using ResultToString for a general case. - static std::string SingleFailureResultToString( - const AllocationResult& result); - // Return a string representation of the result, with possibly more than one - // failure. - static std::string ResultToString(const AllocationResult& result); - - // Return true if the result is a failure that requires us to uncommit pending - // chunks. - static bool result_requires_uncommit(AllocationResult result) { - return result_is(result, AllocationResult::kFailRequiresUncommit); - } - - // Return true if the result is a failure either due to running out of - // outstanding asynchronous copies or due to violating asynchronous copy - // ordering. - static bool result_failed_because_of_async_copy(AllocationResult result) { - return result_is(result, AllocationResult::kFailOutOfAsyncCopies) || - result_is(result, AllocationResult::kFailViolatesAsyncCopyResource); - } - - // Converts an std::optional to a string for - // logging. - static std::string OptionalRequiredMemoryAssignmentToString( - const std::optional& assignment); - - // For the given loop with the start and end index and loop size, run the - // MemoryBoundLoopOptimizer and record its outputs into - // optimized_allocations_map_. - absl::Status OptimizeMemoryBoundLoop(int loop_start_idx, int loop_end_idx, - int loop_size); - - // Identify memory-bound loops in the graph and call OptimizeMemoryBoundLoop - // for the found loops. - void IdentifyAndOptimizeMemoryBoundLoops(); - - // Returns true if the instruction meets the preconditions of a replaceable - // synchronous copy or slice instruction. This only checks for necessary - // conditions, and doesn't guarantee a successful replacement. - bool IsAsyncConversionCandidate(const HloInstruction* instruction) const; - // Not supported instructions for sync copy replacement: - // 1. Layout-changing copies - // 2. Instruction operand or output has a pre-specified memory space - bool IsAsyncConversionCopyCandidate(const HloInstruction* instruction) const; - - enum class AsyncConversionResult { - kSuccess = 0, - kFeatureNotEnabled = 1, - kFailedPrecondition = 2, - kFailedValueNotAllowedInAlternateMemory = 4, - kFailedSatisfyingConstraints = 8, - kFailedNotProcessed = 16, - kFailedGaveUp = 32, - kAsyncConversionNotAllowedForColoredBuffer = 64, - kSourceBufferInAlternateMemory = 128, - }; - - AsyncConversionResult IsAsyncCustomFusionConversionCandidate( - const HloInstruction* instruction) const; - - AsyncConversionResult IsAsyncConversionSliceCandidate( - const HloInstruction* instruction) const; - - // Allocates buffers for instructions that need reserved scoped allocations in - // the alternate memory space. - void AllocateReservedScopedAllocations(); - - // Creates a ScopedAllocation for the given instruction. - void AllocateScopedAllocation(HloInstruction* instruction, - bool is_post_module, int64_t size, int64_t time, - std::vector& colocations); - - // Returns the AliasedOffset object associated with the allocation. - AliasedOffset* GetAliasedOffset(const Allocation& allocation); - - // If aliased_offset is non-null, this method adds the allocation to - // aliased_offset. Otherwise, it creates a new AliasedOffset object and adds - // the allocation to this new AliasedOffset. - void MaybeCreateOrAddToAliasedOffset(const Allocation& allocation, - AliasedOffset* aliased_offset); - - // Records the aliased offset for async pipelined while loop boundary buffers - // to ensure exact colocation across loop parameters, body roots, and callers. - void RecordAliasedOffsetForAsyncPipelinedWhileLoop( - const HloPosition& position, AliasedOffset* aliased_offset); - - // Given an allocation sequence, returns the live allocation at time with a - // preference towards allocations in alternate memory. Returns nullptr if no - // allocation is alive at that time. - static Allocation* GetLiveAllocationAt(const AllocationSequence& allocations, - int64_t time); - - // Returns true if the use is allowed in the alternate memory (hard - // constraints). - bool IsUseAllowedInAlternateMemory(const HloUse& use) const; - - // Returns true if the use is allowed in alternate memory for while loops. - bool IsWhileLoopUseRequiredInDefaultMemory(const HloUse& use) const; - - // Returns true if it is beneficial to have the while loop use in alternate - // memory. - bool IsWhileLoopUseBeneficialInAlternateMemory(const HloUse& use) const; - - // Returns true if it is beneficial to have the conditional use in alternate - // memory. - bool IsConditionalUseBeneficialInAlternateMemory(const AllocationValue& value, - const HloUse& use) const; - - // Adjusts the preferred memory offset for a given use, taking aliasing - // constraints into account. If the use already has a preferred offset in the - // alternate memory space (e.g., due to prior allocations), the offset derived - // from aliasing considerations must match the existing preferred offset. - AliasedOffset* CheckOrUpdatePreferredOffsetForUse( - const AllocationValue::Use& use, AliasedOffset* preferred_offset) const; - - // Propagate the allocation at the use time to any aliases that this use might - // have had. Returns false without recording anything when the propagated - // requirement contradicts an already recorded required assignment (or a - // buffer coloring) at an aliased position: the caller should treat the - // use's allocation as failed and uncommit/retry rather than proceed into a - // CHECK failure inside AddRequiredAssignment. - [[nodiscard]] bool TryUpdateAllocationRequirementForUseAliases( - const AllocationValue& allocation_value, const AllocationValue::Use& use, - int64_t use_time); - - // For while uses that are allocated in the alternate memory space, if - // they also have an allocation in the default memory space in their - // allocation sequence, create a "parent" allocation that mirrors this - // default memory space allocation. When we process the parent - // allocation, we add an additional parameter to the while that is a - // reference to the buffer in the default memory space. With parent - // allocations, we don't need to unnecessarily evict buffers since they - // already have a copy in the default memory space. We search backwards - // (latest to earliest in execution time) for a suitable allocation in - // order to find the most recent one. - // - // Updates the following data structures: - // - `allocation_values`: Appends `ParentAllocation` and `MirroredAllocation` - // entries to the `AllocationSequence`s of the matching while body parameter - // and post-while `AllocationValue`s if redundant eviction elimination is - // enabled. - // - `preferred_offset_for_computation`: Maps the while body computation to - // the preferred `AliasedOffset*` when `has_async_pipelined_while_loops_` is - // false. - // - Delegates to `SynchronizeAliasedWhileLoopOffsets` when - // `has_async_pipelined_while_loops_` is true. - void MaybeCreateMirroredParentAllocationForWhileUse( - const AllocationValue& allocation_value, const AllocationValue::Use& use, - int64_t use_time, absl::Span allocation_values, - absl::flat_hash_map& - preferred_offset_for_computation); - - // Synchronizes the preferred alternate memory offset and records required - // memory assignments across all aliased positions of an async pipelined while - // loop (the loop parameter, body root instruction, and while instruction - // itself) for a given tuple element (`hlo_use.operand_index`). - // - // Updates the following member data structures: - // - `pipelined_while_preferred_offset_for_computation_`: Maps - // `(while_body, operand_index)` to `offset`. - // - `required_assignments_`: Adds required `kAliasedUse` assignments on the - // while body parameter, while body root, and while instruction via - // `AddAliasedRequiredAssignment`. - // - `pipelined_while_buffer_id_to_aliased_offset_`: Maps the - // `HloBuffer::id()` of the while body parameter, while body root, and while - // instruction positions to `offset` via - // `RecordAliasedOffsetForAsyncPipelinedWhileLoop`. - void SynchronizeAliasedWhileLoopOffsets(const HloUse& hlo_use, - const Allocation& aliased_allocation, - AliasedOffset* offset); - - // Returns true if an `allocation_value` satisfies a set of conditions - // indicating that nested allocation values should reuse the alternate memory - // allocated for `allocation_value` between `previous_use` and `current_use`, - // by deploying a MirroredAllocation. - bool ShouldBeMirrored( - const AllocationValue& allocation_value, - const AllocationValue::Use& current_use, - // We check if the previous use is a conditional operand. - const AllocationValue::Use* previous_use) const; - - // Does the work of ShouldBeMirrored for nested conditionals. - bool ShouldBeMirroredByNestedConditionals( - const AllocationValue& allocation_value, - const AllocationValue::Use& current_use, - const AllocationValue::Use* previous_use) const; - - // If `allocation_value` ShouldBeMirrored, create all necessary - // MirroredAllocations. - void CreateMirroredAllocations( - AllocationValue& allocation_value, - const AllocationValue::Use& current_use, - // We check if the previous use is a conditional operand. - const AllocationValue::Use* previous_use, - absl::Span allocation_values, - absl::flat_hash_set& - already_processed_allocation_values_inside_a_conditional); - - // Does the work of CreateMirroredAllocations for nested conditionals. - void CreateMirroredAllocationsForNestedConditionals( - AllocationValue& allocation_value, - const AllocationValue::Use& current_use, - const AllocationValue::Use* previous_use, - absl::Span allocation_values, - // A set of allocation values inside the conditional, that may get a - // mirrored allocation that points to a real allocation outside the - // conditional, that is live throughout the conditional. We maintain - // this set to avoid re-processing these allocation values. - absl::flat_hash_set& - already_processed_allocation_values_inside_a_conditional); - - // Returns true, if the previous use is a conditional operand in the alternate - // memory, and, an eviction is required before the conditional. We check if - // all the buffer positions and uses inside the conditional are allowed in - // alternate memory and if the jointly processed allocation values can be - // processed without imposing infeasible constraints. We require an eviction - // if these conditions are not met. - bool IsEvictionRequiredForPreviousUseAtConditional( - AllocationValue& allocation_value, const AllocationValue::Use& use, - const AllocationValue::Use* previous_use, - absl::Span allocation_values); - - // Creates a detailed memory allocation request for a given use of an - // allocation value. Analyzes the usage pattern of the use to determine if it - // can be placed in alternate memory, considering the restrictions for loops - // and conditionals. Also calculates the timing for prefetching, taking into - // account instruction schedules, operation type (e.g., sequential vs. - // non-sequential calls), and prior usage patterns. We add the resulting - // Allocation to the AllocationSequence of allocation_value_to_update. When - // only_extend_existing_allocation is true, no new Allocations will be created - // while processing the resulting AllocationRequest, and we only need to - // extend an existing Allocation's end_time. - // - // * processed_allocation_values: The AllocationValues that have already been - // processed for the same parent HloValue as is used in the request. - // * shape_override: This shape if provided will be used to determine the - // space needed for the allocation. It overrides the shape used in - // PrefetchContext. - AllocationRequest CreateAllocationRequest( - AllocationValue& allocation_value, - AllocationValue& allocation_value_to_update, - const AllocationValue::Use& use, const AllocationValue::Use* previous_use, - AliasedOffset* preferred_offset, int64_t definition_time, - bool require_no_copy_alternate_mem_allocation, - const std::vector& all_use_times, - bool only_extend_existing_allocation, - absl::Span processed_allocation_values, - absl::Span all_allocation_values, - std::optional shape_override); - - // Returns true, if the allocation value requires a pinned allocation in the - // alternate memory space. - bool RequiresNoCopyAlternateMemAllocation( - AllocationValue& allocation_value) const; - - // Adds a required assignment in default memory, at the given time, if - // allocation_value's position is not allowed in alternate memory. - void AssignDefaultMemIfNotAllowedInAlternateMem( - AllocationValue& allocation_value, int64_t time); - - // Returns all AllocationSegmentContexts needed for a given set of - // AllocationValues that we would like to process jointly. - std::vector GenerateAllocationSegmentContexts( - absl::Span& allocation_values, - absl::flat_hash_map>& - value_indices_by_sync_inst, - int allocation_value_idx) const; - - bool VerifyAllConversionsAreSuccessful(); - - // For async pipelined while loops, updates and returns whether the given - // allocation value should require a pinned (no-copy) allocation in alternate - // memory space. - // - // Returns false if: - // - `preferred_offset` is null. - // - The defining position is not allowed in alternate memory or exceeds - // `max_size_in_bytes`. - // - The buffer is not aliased to an async pipelined while loop. - // - The defining instruction is an asynchronous operation other than dynamic - // update slice or dynamic slice. - // - The underlying buffer ID is in `default_req_buffer_ids` (a set of buffer - // IDs that have required assignments in default memory). - // Otherwise, returns true. - bool GetUpdatedRequireNoCopyAlternateMemForAsyncPipelinedWhile( - bool require_no_copy_alt_mem, - const AllocationValue& allocation_value_to_update, - const AliasedOffset* preferred_offset, - const absl::flat_hash_set& default_req_buffer_ids) const; - - // Finds allocations for allocation values generated from colocated intervals. - // All of the allocation values have a must-alias relationship with each - // other. Returns either kSuccess if all of the sites could be placed in the - // alternate memory or a bitwise OR of failure reasons why they couldn't - absl::StatusOr AllocateAllocationValues( - absl::Span allocation_values); - - // Returns true if the conditional outputs should be required in the default - // memory. Reasons include: - // * There are required assignments in the default memory for the conditional - // output position or any of its aliases. - // * At least one of the branched computation root instructions is not a tuple - // instruction. This is due to a limitation of MSA, where in, an - // AllocationValue can only be assigned VMem if it has at least one HloUse. - // In this case, we make sure that the HloValue is not defined at the root - // instruction and has at least one use - the tuple instruction. - bool RequireConditionalOutputsInDefaultMemory( - HloPosition conditional_phi_position, const HloValue* hlo_value); - - // Checks for a situation in which an HloValue has more than one live - // AllocationValue at the same time, and the already processed AllocationValue - // has been given alternate memory at the start of the second AllocationValue. - // If such a case is detected, we set - // request.no_copy_chunk_inclusive_start_time with the time where the first - // AllocationValue left off. AllocateInAlternateMemoryNoCopy() takes advantage - // of that information. - void CheckAndUpdateForDualLiveAllocationValues( - const std::optional& - required_memory_assignment_at_start, - AllocationRequest& request); - - // Finds an allocation for an allocation request for a segment (see the - // documentation for AllocationRequest above how a segment is defined). - // - // It performs three things in the following order: - // 1- Allocate the allocation request entirely in the alternate memory, if - // there is enough space and if the prefetch interval picker allows. - // 2- If (1) was unsuccessful, and the only allocation for - // this buffer was in the alternate memory, we try to perform a prefetch. - // 3- If (1) was unsuccessful, prefetch the buffer into the alternate memory, - // if there is enough space and if the prefetch interval picker allows. - // - // If an eviction (2) was requested and was unsuccessful, this method returns - // Result::kFailRequiresUncommit. This means we could not find a suitable - // allocation, so all previous allocations for this buffer must be removed and - // allocated in the default memory. Otherwise, this method may return - // Result::kSuccess if the buffer could be placed in alternate memory or some - // other Result with an OR of reasons why the buffer couldn't be placed in - // alternate memory. - AllocationResult AllocateSegment(AllocationRequest& request); - - // Try allocating in alternate memory without any copies. - AllocationResult AllocateInAlternateMemoryNoCopy( - const AllocationRequest& request); - - // Try allocating in alternate memory for the minimum time possible. - AllocationResult ForceAlternateMemoryAllocationForMinTime( - const AllocationRequest& request); - - // Try evicting to default memory space. If force_evict is true, we will - // evict even if the resource constraints for an eviction are not met. - AllocationResult Evict(const AllocationRequest& request, - bool force_evict = false); - - // Returns the time a copy done of a prefetch should be scheduled. - int64_t FindPrefetchEndTime(const AllocationRequest& request, - int64_t earliest_prefetch_time) const; - - // Try prefetching to alternate memory space. If force_prefetch is true, we - // will prefetch even if the resource constraints for a prefetch are not met. - AllocationResult Prefetch(const AllocationRequest& request, - Allocation& prev_allocation_in_default_mem, - bool force_prefetch = false); - - // Prefetch to alternate memory iff the resource constraints are met. - AllocationResult PrefetchWithResourceConstraints( - const AllocationRequest& request, - Allocation& prev_allocation_in_default_mem); - - // Helper methods used to implement Prefetch(). - // - // Generates a SliceProposal in context, if options dictate and one can be - // constructed. - void GenerateSliceProposal(PrefetchContext& context) const; - // Calls GenerateSliceProposal to potentially create a SliceProposal, and - // sets up WorkingIntervals for a sliced and unsliced solution. Updates - // context. - void SetupPrefetchWorkingIntervalsAndSliceProposal( - PrefetchContext& context) const; - // Initializes the PrefetchIntervalPicker and associated data structures in - // context. - AllocationResult InitializePrefetchIntervalPicker(PrefetchContext& context); - // As a compile time optimization, try a prefetch allocation that is as late - // as possible. If this is not able to find a solution, none of the - // earlier tries will succeed either. - AllocationResult EnsureSomeSpatialPrefetchFitExists( - PrefetchContext& context) const; - // Check if for the specified type of solution, using the parameters in - // context. If we find a solution, it will be stored in context. - AllocationResult CheckPrefetchFit(bool for_sliced_solution, - PrefetchContext& context); - // Creates a debugging string describing the timing of the prefetch solution - // we are currently attempting (as dictated by for_sliced_solution and - // context). - std::string AlternateMemoryAllocationAttemptToString( - bool for_sliced_solution, const PrefetchContext& context) const; - - // Performs window prefetching. - absl::Status WindowPrefetch(); - - // Window prefetches the specified operand of the given instruction. - void WindowPrefetchOperand(const HloUse& use, int64_t bytes); - - // Find the best possible chunk candidate, where it has the longest - // possible availability if no preferred offset is given, or at the - // preferred_offset if it is given. - std::optional FindBestChunkCandidate( - const AllocationRequest& request, const AliasedOffset* preferred_offset, - MsaBufferInterval* alternate_mem_interval) const; - - // Returns the corrected schedule time of an HloUse. The corrected time is - // equivalent to the actual time of the use instructions for all instructions - // except for while and conditional instructions. For while instructions, the - // corrected time is the time of the body parameter, and for conditional, the - // corrected time is the time of the parameter of the earliest-scheduled - // called computation. - int64_t GetCorrectedUseTime(const HloUse& use) const; - int64_t GetCorrectedUseTime(const HloInstruction* instruction) const; - - // If `use` is a view, returns the extended use time for the pointed-to - // allocation; otherwise returns the time of the use. - int64_t GetExtendedUseTimeIfUseIsView(const HloUse& use) const; - - // Returns the required assignment at a particular time, if available. - std::optional RequiredMemoryAssignmentAt( - const HloValue* buffer, int64_t time) const; - - // Returns true if requiring `sites` in default memory would record at least - // one requirement that is not already in effect. Pending requirements do not - // count: UncommitPendingWork rolls them back before the retry. Returns false - // when every site already carries a committed (non-pending) required - // assignment, in which case the retry would be a no-op: it would run against - // exactly the same state, produce the same allocations, and flag the same - // sites again. - bool InefficientSiteRetryCanProgress( - absl::Span sites) const; - - // Searches for aliases in the use for a required assignment, and returns it - // if found. - std::optional AliasedRequiredAssignmentForUse( - const AllocationValue::Use& use) const; - - // Returns the required assignment for a given use. - std::optional RequiredAssignmentForUse( - const AllocationValue::Use& use) const; - - // Adds required assignment in the default memory for conditional outputs - // that RequireConditionalOutputsInDefaultMemory() returns true for. - void AddRequiredAssignmentsForConditionalOutputsIfNecessary( - absl::Span colocated_intervals); - - // Propagates aliased required assignment for a given position. - void AddAliasedRequiredAssignment( - const HloInstruction* instruction, ShapeIndex index, - const Allocation* aliased_allocation, - RequiredMemoryAssignment::Source required_assignment_source); - - // This sets a required assignment. CHECK fails if there is a conflicting - // required assignment at the same time. - void AddRequiredAssignment( - const HloValue* value, const HloInstruction* instruction, - MemorySpace memory_space, int64_t time, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset = nullptr, bool add_to_pending = true); - void AddRequiredAssignment( - const HloInstruction* instruction, ShapeIndex index, - MemorySpace memory_space, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset = nullptr, bool add_to_pending = true); - void AddRequiredAssignment( - const HloPosition& position, MemorySpace memory_space, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset = nullptr, bool add_to_pending = true); - void AddRequiredAssignment( - const HloUse& use, MemorySpace memory_space, - RequiredMemoryAssignment::Source required_assignment_source, - AliasedOffset* offset = nullptr, bool add_to_pending = true); - - // Adds input and outputs as required assignments. - void AddInputAndOutputRequiredAssignments(); - - // Returns a list of "linked" allocations in the alternate memory. Linked - // allocations all share a common allocation site (a use or position) with - // each other. This can be used to determine if a group of linked allocations - // are considered efficient or not. - std::vector> - GetLinkedAllocationsInAlternateMemory( - absl::Span allocation_values) const; - - // Returns allocation sites (use or position) that are allocated in the - // alternate memory, but is considered inefficient. These arise in the - // context of in-place operation like dynamic-update-slice. We will typically - // have an allocation that has the DUS as a use, and another allocation that - // has the DUS as a defining position. These two allocation will be part of - // the same linked allocation group. - // - // One reason why an allocation site could be inefficient is because the - // amount of data that is asynchronously copied (prefetch and eviction) is - // much larger than the amount of data that is used by the HLOs. If we find - // inefficient allocation sites, we can require these sites default memory - // allocations and allocate them again. - std::vector GetInefficientAllocationSites( - absl::Span allocation_values) const; - - // Returns true if the colocated intervals in the argument are in a parameter - // or root instruction of the entry computation and are reserved by the user - // to be in the alternate memory space. - bool AreIntervalsReservedInAlternateMemory( - absl::Span colocated_intervals) const; - - // Since the allocations are recorded to the AllocationSequence, we don't - // maintain result_ in GlobalDecreasingSizeBestFitHeap. Override AddToChunkMap - // to avoid unnecessarily adding the chunk to the chunk map. - // - // Sliced prefetching requires that we override this method because we - // associate more than one chunk with a buffer (i.e., 1 chunk per slice), - // which would cause the original implementation of this method to CHECK fail. - void AddToChunkMap(const HloValue* buffer, Chunk chunk) override {} - - // Returns true if the addition of num_additional_copies asynchronous copies - // in the given time interval would violate the maximum number of asynchronous - // copies. An extra async copy limit can be provided to increase the limit of - // asynchronous copies for this instance. - bool ViolatesMaximumOutstandingAsyncCopies( - int64_t inclusive_start_time, int64_t end_time, bool is_prefetch, - int64_t extra_async_copy_limit = 0, - int64_t num_additional_copies = 1) const; - - // Exports the allocations for repacking and puts them into the vector in the - // parameter. - void ExportAllocationsForRepacking( - std::vector& allocations); - - // Update reserved scoped allocation size for instructions when their - // operand/output has been allocated in alternate memory by invoking - // reserved_scoped_memory_fn - void UpdateReservedScopedAllocationSize(); - - // Uncommits the chunk and updates the peak memory. This method keeps other - // data structures consistent with the uncommit. - // - // MSA should not directly update interval_tree_ or directly call other - // Uncommit methods. This the definitive uncommit for MSA. - bool UncommitChunkAndUpdatePeakMemory(const MsaBufferInterval& interval, - const Chunk& chunk); - - // Commits the chunk and updates the peak memory. This method keeps other - // data structures consistent with the commit. - // - // MSA should not directly update interval_tree_ or directly call other - // Commit methods, e.g., CommitChunkAndInterval(). This the definitive - // uncommit for MSA. - void CommitChunkAndUpdatePeakMemory(const MsaBufferInterval& buffer_interval, - const Chunk& chunk); - - // Imports repacked allocations and updates the internal data structures - // consistent with the new packing. - void ImportRepackedAllocations(); - // Helper functions to implement ImportRepackedAllocations. - void ImportRepackedNonSlicedAllocation(RepackAllocationBlock& block); - void ImportRepackedSlicedAllocation(RepackAllocationBlock& block); - absl::Status AreRepackedSlicesValid(const RepackAllocationBlock& block); - - // Registers an asynchronous copy with asynchronous copy data structures to - // keep track of its state. - void RegisterAsyncCopy(MemorySpace memory_space, int64_t exclusive_start_time, - int64_t copy_done_schedule_before_time, - AllocationSequence* allocations, - AliasedOffset* aliased_offset, float resource, - std::optional cross_program_prefetch_index); - - // Adds an asynchronous copy or other memory operation (e.g., slice) to - // allocations. We pass sync_mem_op to the CopyAllocation constructor. When - // sync_mem_op is set, instead of an async copy, CopyAllocation::Process() - // will replace sync_mem_op with the async version of sync_mem_op's opcode - // (e.g., slice) and shape. Generally, MSA inserts and schedules new async - // copy instructions. If async copy instructions are already present in the - // original schedule, MSA will just schedule them in correct positions. If not - // null, `async_mem_op_start` and `async_mem_op_done` are async copy start and - // done instructions that are already present in the original schedule. - void AddAsyncCopyOrOtherMemOp( - Allocation& prev_allocation, MemorySpace memory_space, - std::optional chunk, int64_t exclusive_start_time, - int64_t end_time, int64_t copy_done_schedule_before_time, - AllocationSequence* allocations, AliasedOffset* aliased_offset, - float resource, - std::optional cross_program_prefetch_index = std::nullopt, - HloInstruction* sync_mem_op = nullptr, - HloInstruction* async_mem_op_start = nullptr, - HloInstruction* async_mem_op_done = nullptr, - int64_t source_operand_index = 0); - - // For prefetching, adds a SlicedCopyAllocation to allocations. Also updates - // asynchronous copy data structures, prefetch_interval_tree_, and aliasing - // data structures - void AddAsyncSlicesForPrefetch( - const Allocation& prev_allocation, AllocationSequence* allocations, - AliasedOffset* aliased_offset, - const std::vector& slice_decisions_sorted_by_start_time, - int64_t prefetch_end_time, int64_t allocation_end_time, - HloInstruction* sync_mem_op); - - // For window prefetching, adds a WindowPrefetchedAllocation to allocations. - // Also updates asynchronous copy data structures, prefetch_interval_tree_, - // and aliasing data structures. - void AddAsyncCopyForWindowPrefetch( - Allocation& prev_allocation, HloUse use, const Chunk& chunk, - int64_t exclusive_start_time, int64_t inclusive_end_time, - AllocationSequence* allocations, AliasedOffset* aliased_offset, - float resource, const WindowPrefetchedAllocation::Options& options); - - // This method is used for committing the chunk candidate but adding it to - // pending_chunks_ so that we can "uncommit" them in case we need to roll back - // this allocation sequence. - void AddToPendingChunks(const MsaBufferInterval& buffer_interval, - const Chunk& chunk); - - // If we need to remove the allocations for this allocation sequence, this - // removes pending chunks and asynchronous copies in the respective pending - // buffers from the interval trees. If an allocation request returns - // kFailRequiresUncommit, this method must be called. - void UncommitPendingWork(absl::Span allocation_values); - - // Finalizes the allocations where they can no longer be uncommitted. - void FinalizeAllocations(absl::Span allocation_values); - - // Clears all pending chunks and asynchronous copies. - void ClearPendingChunks(); - - // Returns true if we are trying to replace instruction with its async - // version, while processing JointAllocationProposal. - bool IsInstructionPendingReplacements( - const HloInstruction* instruction) const; - - // Colors the colocated intervals in the alternate memory. - void ColorColocatedIntervalsToAlternate( - const std::vector& colocated_intervals); - - // A proposal for a group of values to be allocated jointly. Proposals are not - // guaranteed to be accepted, and when they fail, the algorithm will try to - // come up with a new proposal on a smaller subset of values. - struct JointAllocationProposal { - // The values that are being jointly processed. - std::vector values; - // The allocation values created for the joint-processed values. - std::vector allocation_values; - // The colocated buffer intervals for the joint-processed values. This is a - // vector of vectors, one vector per joint-processed value, and the - // colocation must be only enforced on intervals belonging to the same - // joint-processed value. - std::vector> colocated_intervals; - }; - - // Iterates over proposal's values and populates its allocation_values and - // colocated_intervals with the appropriate allocation values and colocated - // intervals created for the values. - void CreateAllocationValuesForJointProcessedValues( - JointAllocationProposal& proposal); - - // Returns a JointAllocationProposal with values, allocation - // values, and colocated intervals that are proposed to be processed jointly - // for the given interval. Also, if the interval consumes or produces any - // synchronous memory call instructions (e.g., kCopy, kSlice) and the option - // to replace them with their asynchronous versions is enabled, this method - // will add those instructions to the sorted_async_conversion_candidates_ - // vector. - JointAllocationProposal GetJointProposal(MsaBufferInterval& interval); - - // Append buffer and allocation infos for debugging and dump it into a file, - // if enabled. - void AppendBufferInfoDebugString(const MsaBufferInterval& interval, - std::string* debug_str) const; - void AppendScopedAllocationBufferInfoDebugString( - const HloInstruction* instruction, int64_t time, int64_t size, - std::string& debug_str) const; - void AppendAllocationInfoDebugString(const Allocation& allocation, - std::string& debug_str) const; - void DumpDebugStringsIfEnabled() const; - - // Returns the available heap size in the alternate memory. - int64_t available_heap_size() const { - return options_.max_size_in_bytes - reserved_in_bytes_; - } - - // Returns the earliest time in the (exclusive_start_time, end_time) range - // that a new allocation with the given size would fit in the alternate - // memory. If it doesn't fit, it returns nullopt. - std::optional FindEarliestExclusiveTimeToSatisfyPeakMemory( - int exclusive_start_time, int end_time, int64_t size) const; - - // Creates and returns a RepackAllocationBlock. - RepackAllocationBlock MakeRepackAllocationBlock(int64_t start_time, - int64_t end_time, - int64_t size, - int64_t initial_offset, - Allocation* allocation) { - RepackAllocationBlock allocation_block; - allocation_block.inclusive_start_time = start_time; - allocation_block.end_time = end_time; - allocation_block.size = size; - allocation_block.offset = -1; - allocation_block.initial_offset = initial_offset; - allocation_block.id = next_repack_allocation_block_id_++; - allocation_block.next_colocated = nullptr; - allocation_block.allocation = allocation; - return allocation_block; - } - - // Returns a vector of instructions that have the same fingerprint as this - // instruction. - const std::vector* GetRepeatedInstructionList( - const HloInstruction* instruction) const; - - // Adds an operand to the alternate memory map. - void AddOperandToAlternateMemoryMap(const HloInstruction* instruction, - int operand_number, - const ShapeIndex& index); - - // Returns true if the interval is pinned in the alternate memory. Buffers are - // pinned when their layout has the alternate memory space before MSA runs. - bool IsIntervalPinnedToAlternateMemory( - const MsaBufferInterval& interval) const; - - // Returns true if the buffer is aliased to the program output. If so, we - // cannot prefetch the buffer into Vmem. - bool IsBufferAliasedToProgramOutput(const HloBuffer* buffer) const; - - // Returns true if the value has at least one use which is not an async - // conversion candidate. In other words, a false return value implies there - // are no uses that would benefit from a copy of the HloValue in Vmem. - bool DoesValueHaveNonAsyncConversionCandidateUses( - const HloValue* value) const; - - // A convenience debugging method that returns true if the prefetch context - // matches the described producer and consumer. - bool MatchesPrefetchContext(const PrefetchContext& context, - absl::string_view producer_name, - ShapeIndex producer_shape_index, - absl::string_view consumer_name) const; - - // Finds the matching AllocationValue for a given HloUse. Returns nullptr if - // no matching AllocationValue is found. - // - // candidate_allocation_values contains all AllocationValue objects for the - // non-trivial defining positions of the specific HloValue that `use` - // consumes. - // - // REQUIRES: candidate_allocation_values must be sorted by the definition time - // of their defining instruction. - AllocationValue* FindAllocationValueForUse( - const HloUse& use, - absl::Span candidate_allocation_values, - int64_t use_time) const; - - // Takes a group of allocation values and splits them if they can be split on - // the same dimension. - void MaybeSplitAllocationValues( - absl::Span allocation_values); - - // Processes the buffer uses that have been colored. Note: Defining position - // of a buffer is also considered as a use that can be colored. - absl::Status ProcessColoredBuffers(); - - // Returns a sorted list of time intervals for which a buffer needs to be - // contiguously allocated. - absl::StatusOr> GetContiguousLiveRangesForBuffer( - const HloBuffer* buffer) const; - - // Process the buffer colorings for the given buffer and return the time - // intervals that are colored in the alternate memory space. Additionally, - // update the default_memory_coloring_requirements_ with the default memory - // colored intervals for the buffer. - absl::StatusOr> - GetAltMemoryColoredIntervalsForBuffer( - const HloBuffer* buffer, - const std::vector& buffer_colorings); - - // Checks if there are any overlapping time intervals between default memory - // colorings and alternate memory colorings for the given buffer. Returns an - // error if conflicts exist. - absl::Status CheckForConflictingColoringRequirements( - const HloBuffer* buffer, - const std::vector& default_mem_intervals, - const std::vector& alternate_mem_intervals) const; - - // Groups the buffer colorings by the HloBuffer that is colored. - absl::linked_hash_map> - GetHloBufferToColoringsMap() const; - - // Removes the reserved chunk from the interval_tree_ for the given - // allocation (if it is still reserved) and removes the corresponding - // RepackAllocationBlock from repack_allocation_blocks_. - void ReleaseReservedAllocationForAlternateMemoryColorings( - ReservedAllocation* allocation); - - // Frees the reserved allocations that are used to satisfy alternate memory - // coloring requirements, for the given allocation request. - void FreeAlternateMemoryColoringReservedAllocations( - AllocationRequest& request); - - // Sets the alternate memory coloring requirements for the given allocation - // request. - void UpdateRequestWithAlternateMemoryColoringRequirements( - AllocationRequest& request); - - // Sets the default memory coloring requirements for the given allocation - // request. - void UpdateRequestWithDefaultMemoryColoringRequirements( - AllocationRequest& request); - - // Whether an hlo position is colored in alternate memory. - bool IsPositionColoredInAlternateMemory(const HloPosition& position) const; - - // Whether an hlo use is colored in alternate memory. - bool IsUseColoredInAlternateMemory(const HloUse& use) const; - - // Whether an hlo position is colored in default memory. - bool IsPositionColoredInDefaultMemory(const HloPosition& position) const; - - // Whether an hlo use is colored in default memory. - bool IsUseColoredInDefaultMemory(const HloUse& use) const; - - // Whether an hlo position is colored in alternate memory at the given time. - bool IsPositionColoredInAlternateMemoryAtTime(const HloPosition& position, - int64_t time) const; - - // Whether an hlo position is colored in default memory at the given time. - bool IsPositionColoredInDefaultMemoryAtTime(const HloPosition& position, - int64_t time) const; - - // Reserves a chunk in alternate memory of size MaxScopedMemorySize() for - // the entire program duration for scoped memory allocations. - int64_t ReserveAlternateMemoryForScopedMemoryAllocations(); - - // Frees the alternate memory reserved for scoped memory allocations. - void FreeAlternateMemoryForScopedMemoryAllocations( - int64_t max_scoped_memory_size); - - // Returns sorted buffer intervals after some custom re-ordering. - std::vector GetPostProcessedSortedBufferIntervals(); - - // Calculates the memory pressure for the buffers that can be assigned in - // alternate memory. - void CalculateMemoryPressure( - const std::vector& sorted_buffer_intervals); - - // Marks the given HloValue as finalized, i.e., its allocation values have - // been finalized and cannot be uncommitted or changed. - void FinalizeValue(const HloValue* value); - - // Returns true if the given HloValue is finalized, i.e., its allocation - // values have been finalized and cannot be uncommitted or changed. - bool IsValueFinalized(const HloValue* value) const { - return finalized_values_.contains(value); - } - - bool IsBlockPrefetchingEnabled() const; - - HloModule* module_ = nullptr; - AllocationSequence* allocations_; - // Edge time indices store start and end times allocations in alternate - // memory. We can use this to skip redundant calls to FindChunkCandidate. - absl::flat_hash_set edge_time_indices_; - const Options& options_; - const HloAliasAnalysis& alias_analysis_; - const AliasInfo* alias_info_; - const HloLiveRange& hlo_live_range_; - std::unique_ptr call_graph_; - // We use a interval tree to keep track of the number of outstanding - // prefetches and evictions. - BufferIntervalTree prefetch_interval_tree_; - BufferIntervalTree eviction_interval_tree_; - AsynchronousCopyOrdering async_copy_ordering_; - AsynchronousCopyResource prefetch_async_copy_resource_; - AsynchronousCopyResource eviction_async_copy_resource_; - // A list of RepackAllocationBlock objects that mirrors allocation sequences, - // used for repacking. We use a list here because we need pointer stability - // for aliased allocations. - std::list repack_allocation_blocks_; - int64_t next_repack_allocation_block_id_ = 0; - int64_t num_repacks_ = 0; - int64_t num_repacks_successful_ = 0; - // True if any instruction in the module carries frontend attributes - // indicating custom while loop copy control, requiring async pipelined - // while loop alternate memory offset colocation. - bool has_async_pipelined_while_loops_ = false; - std::vector> pending_chunks_; - std::vector pending_async_copies_; - std::vector> - pending_required_assignments_; - // A list of candidate sync instructions that we are trying to replace with - // an asynchronous version, while processing the current interval, sorted by - // their order in the instruction schedule. Being in this list doesn't - // guarantee that the sync instruction will be converted to async. - std::vector sorted_async_conversion_candidates_; - // A cache to keep the peak memory usage at each point in the graph. We use - // this to see if the proposed allocation in the alternate memory would fit - // ignoring fragmentation, and if not, we can skip the more expensive lookup - // in the BufferIntervalTree, which also considers fragmentation. - std::vector peak_memory_usage_; - // The data structure that contains AliasedOffset objects and Allocation to - // AliasedOffset map for efficient lookup. - std::list aliased_offsets_; - absl::flat_hash_map aliased_offset_map_; - // This map contains required memory assignments for HloValues (e.g., input - // and outputs). - absl::flat_hash_map> - required_assignments_; - // Number of bytes reserved in alternate memory space. - int64_t reserved_in_bytes_ = 0; - // A rough measure of the memory pressure of the model, in bytes. Note that - // this is pressure for memory capacity (and not accessed bytes), and for - // alternate memory (not default memory). - int64_t memory_pressure_ = 0; - int64_t next_async_copy_id_ = 0; - // Fingerprint cache. - absl::flat_hash_map fingerprint_map_; - // Vector of repeated instructions (that have the same fingerprint) indexed by - // fingerprint. - absl::flat_hash_map> - repeated_inst_map_; - - // Loop-optimized allocations found by MemoryBoundLoopOptimizer. These - // allocation objects describe the allocations for one iteration of the loop, - // so we translate them into the program-level Allocation objects in - // allocations_. - std::vector loop_optimized_allocations_; - // A map to look up the loop-optimized allocation info by use. - absl::flat_hash_map - loop_optimized_allocations_map_; - // A map to look the operands of each instruction that are assigned in - // alternate memory or are window prefetched. - absl::flat_hash_map>> - operands_in_alternate_memory_map_; - // A map to look the outputs of each instruction that are assigned in - // alternate memory. - absl::flat_hash_map> - outputs_in_alternate_memory_map_; - // HloValues whose allocation values have been finalized and cannot be - // uncommitted or changed. - absl::flat_hash_set finalized_values_; - // Set of sync copy instructions that we failed/succeeded in replacing with - // asynchronous copies. - absl::flat_hash_map - failed_async_conversions_; - absl::flat_hash_set successful_async_conversion_set_; - std::vector not_finalized_async_conversions_; - // Maps from an HloValue to the dimension it is split on. - absl::flat_hash_map> - instruction_to_split_dims_; - // Debug strings. - std::string buffer_info_str_; - std::string allocation_info_str_; - std::string instruction_schedule_str_; - - // Maps defining HloPositions to the chunk intervals that are reserved for it - // in alternate memory, in order to satisfy buffer coloring requirements. - absl::flat_hash_map>> - reserved_allocations_for_alt_mem_colorings_; - - // Maps defining HloBuffers to the list of time intervals it is required to - // be in default memory, to meet buffer coloring requirements. - absl::flat_hash_map> - default_memory_coloring_requirements_; - - // Set of HloUses that are in the default memory. - absl::flat_hash_set uses_in_default_memory_set_; - - // Vector to preserve insertion order for deterministic window prefetching - // results. - std::vector uses_in_default_memory_; - - // Preferred alternate memory offset mapping for while loop body computations, - // keyed by the while body computation and tuple shape index. Only used for - // pipelined async while loops. Maintained only during the current allocation - // attempt. - absl::flat_hash_map> - pipelined_while_preferred_offset_for_computation_; - - // Mapping from HloBuffer ID to assigned alternate memory offset to ensure - // exact offset colocation across loop parameters, roots, and callers. Only - // used for pipelined async while loops. Maintained only during the current - // allocation attempt. - absl::flat_hash_map - pipelined_while_buffer_id_to_aliased_offset_; - - // We have released the chunks corresponding to the allocations in the list. - // When we uncommit the current pending state following a - // kFailRequiresUncommit, we need to re-reserve those chunks. - std::vector pending_deallocated_reserved_allocations_; -}; - -// Helper to inspect the async wrapped opcode of a pipelined while loop position -// or tuple index. -std::optional GetAsyncPipelinedWhileWrappedOpcode( - const HloInstruction* while_instr, const HloPosition& pos, - const HloAliasAnalysis& alias_analysis); -std::optional GetAsyncPipelinedWhileWrappedOpcode( - const HloInstruction* while_instr, int64_t tuple_idx, - const HloAliasAnalysis& alias_analysis); - -// Returns true if the position in an async pipelined while loop corresponds to -// a buffer that is intended to reside in alternate memory (e.g., prefetched -// dynamic-slice output, or dynamic-update-slice update slice). Base tensors of -// dynamic-slice or dynamic-update-slice and dynamic-update-slice outputs -// reside in default memory (HBM) on TPU and return false. -bool IsAsyncPipelinedWhileAlternateMemoryPosition( - const HloPosition& pos, const HloAliasAnalysis& alias_analysis); - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALGORITHM_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.cc b/third_party/xla/xla/service/memory_space_assignment/allocation.cc deleted file mode 100644 index 875ecdbb5e62cd..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/allocation.cc +++ /dev/null @@ -1,1402 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/allocation.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/base/casts.h" -#include "absl/container/flat_hash_set.h" -#include "absl/functional/function_ref.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/layout.h" -#include "xla/layout_util.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/service/time_utils.h" -#include "xla/service/tuple_util.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -#include "xla/util.h" -#include "xla/xla_data.pb.h" -#include "tsl/platform/casts.h" -#include "tsl/platform/errors.h" -#include "tsl/platform/statusor.h" - -namespace xla::memory_space_assignment { -namespace { - -std::string UsesToString(const std::vector& uses) { - if (uses.empty()) { - return "none"; - } - std::vector uses_str; - uses_str.reserve(uses.size()); - for (const auto& use : uses) { - uses_str.push_back(use.ToString()); - } - return absl::StrJoin(uses_str, ","); -} - -// Helper function to compute the start time for a SlicedCopyAllocation. -int64_t GetSlicedCopyAllocationExclusiveStartTime( - const std::vector& - slice_decisions_sorted_by_exclusive_start_time) { - if (slice_decisions_sorted_by_exclusive_start_time.empty()) { - return -1; - } - - return slice_decisions_sorted_by_exclusive_start_time.front() - .exclusive_start_time; -} - -// Helper function to compute the underlying Allocation chunk for a -// SlicedCopyAllocation. -std::optional GetSlicedCopyAllocationChunk( - const std::vector& slice_decisions_sorted_by_start_time) { - if (slice_decisions_sorted_by_start_time.empty()) { - return std::nullopt; - } - auto offset_cmp = [](const SliceDecision& lhs, const SliceDecision& rhs) { - return lhs.chunk.offset < rhs.chunk.offset; - }; - auto end_cmp = [](const SliceDecision& lhs, const SliceDecision& rhs) { - return lhs.chunk.chunk_end() < rhs.chunk.chunk_end(); - }; - return HeapSimulator::Chunk::FromOffsetEnd( - std::min_element(slice_decisions_sorted_by_start_time.begin(), - slice_decisions_sorted_by_start_time.end(), offset_cmp) - ->chunk.offset, - std::max_element(slice_decisions_sorted_by_start_time.begin(), - slice_decisions_sorted_by_start_time.end(), end_cmp) - ->chunk.chunk_end()); -} - -} // namespace - -std::string MemorySpaceToString(MemorySpace memory_space) { - switch (memory_space) { - case MemorySpace::kDefault: - return "def"; - case MemorySpace::kAlternate: - return "alt"; - } -} - -std::optional Allocation::cross_program_prefetch_index() const { - return cross_program_prefetch_index_; -} - -HeapSimulator::Chunk Allocation::chunk() const { - CHECK(chunk_.has_value()); - return *chunk_; -} - -void Allocation::set_offset(int64_t offset) { - CHECK(chunk_.has_value()); - *chunk_ = HeapSimulator::Chunk::FromOffsetSize(offset, chunk_->size); -} - -bool Allocation::is_in_alternate_mem() const { - return memory_space_ == MemorySpace::kAlternate; -} - -bool Allocation::is_in_default_mem() const { - return memory_space_ == MemorySpace::kDefault; -} - -void Allocation::RemoveUse(HloUse use) { - uses_.erase(std::remove_if(uses_.begin(), uses_.end(), - [=](const auto& u) { return u == use; }), - uses_.end()); -} - -void Allocation::AddUse(HloUse use) { - HloInstruction* operand = - use.instruction->mutable_operand(use.operand_number); - // If the use is a tuple, look inside the tuple to find the actual use. - for (int64_t index : use.operand_index) { - if (operand->opcode() != HloOpcode::kTuple) { - break; - } - operand = operand->mutable_operand(index); - } - - // Look beyond GetTupleElement(Tuple()) pattern for any bitcasts. - std::function get_simplified_operand; - get_simplified_operand = [&](HloInstruction* instruction) { - while (instruction->opcode() == HloOpcode::kGetTupleElement) { - HloInstruction* operand = - get_simplified_operand(instruction->mutable_operand(0)); - if (operand->opcode() == HloOpcode::kTuple) { - instruction = operand->mutable_operand(instruction->tuple_index()); - } else { - return instruction; - } - } - return instruction; - }; - operand = get_simplified_operand(operand); - - uses_.push_back(use); -} - -absl::Status Allocation::UpdateUses(HloComputation* computation, - HloInstruction* producing_instruction, - const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - for (const HloUse& use : uses()) { - HloInstruction* replacement_instruction = producing_instruction; - const Shape& operand_shape = - use.instruction->operand(use.operand_number)->shape(); - if (operand_shape.IsTuple()) { - HloInstruction* tuple_inst = - use.instruction->mutable_operand(use.operand_number); - - // Trace back to original instructions in the schedule to use - // HloAliasAnalysis safely. - auto find_original_leaf = [&](const HloInstruction* inst, - ShapeIndex index) - -> std::pair { - while (!hlo_live_range.instruction_schedule().contains(inst)) { - if (inst->opcode() == HloOpcode::kTuple) { - if (index.empty()) { - break; - } - int64_t op_idx = index[0]; - inst = inst->operand(op_idx); - index = ShapeIndex(index.begin() + 1, index.end()); - } else if (inst->opcode() == HloOpcode::kGetTupleElement) { - ShapeIndex new_index; - new_index.push_back(inst->tuple_index()); - for (int i : index) { - new_index.push_back(i); - } - inst = inst->operand(0); - index = new_index; - } else if (inst->opcode() == HloOpcode::kBitcast || - inst->opcode() == HloOpcode::kCopy || - inst->opcode() == HloOpcode::kCopyDone || - inst->opcode() == HloOpcode::kCopyStart) { - inst = inst->operand(0); - } else { - break; - } - } - return {inst, index}; - }; - - bool should_skip_reconstruction = false; - - // If the producing instruction shares the same buffer with the operand, - // then we skip the reconstruction. - auto is_aliasing_operand = [&](int operand_index) { - ShapeIndex sub_index; - if (use.operand_index.size() > 1) { - sub_index = ShapeIndex(use.operand_index.begin() + 1, - use.operand_index.end()); - } - auto [traced_operand, traced_index] = - find_original_leaf(tuple_inst->operand(operand_index), sub_index); - return hlo_live_range.instruction_schedule().contains(traced_operand) && - hlo_live_range.instruction_schedule().contains(tuple_inst) && - alias_analysis.GetUniqueBufferAt(traced_operand, traced_index) - .id() == - alias_analysis - .GetUniqueBufferAt(tuple_inst, use.operand_index) - .id(); - }; - - for (int operand_index = 0; operand_index < tuple_inst->operand_count(); - ++operand_index) { - if (tuple_inst->opcode() == HloOpcode::kTuple && - !use.operand_index.empty() && - operand_index != use.operand_index.front()) { - continue; - } - if (tuple_inst->operand(operand_index) == producing_instruction && - is_aliasing_operand(operand_index)) { - should_skip_reconstruction = true; - break; - } - } - - if (should_skip_reconstruction) { - VLOG(4) << "Skipping tuple reconstruction for " << tuple_inst->name() - << " at index " << use.operand_index.ToString(); - replacement_instruction = tuple_inst; - } else { - ABSL_ASSIGN_OR_RETURN( - replacement_instruction, - TupleUtil::ReplaceTupleWith(producing_instruction, tuple_inst, - use.operand_index)); - } - } 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. - VLOG(4) << "Old shape = " << operand_shape.ToString() - << ", new shape = " << producing_instruction->shape().ToString() - << "; inserting a bitcast."; - replacement_instruction = computation->AddInstruction( - HloInstruction::CreateBitcast(operand_shape, producing_instruction)); - if (mutable_split_shape().has_value() && - producing_instruction->shape().layout().split_configs().size() != 0) { - ABSL_ASSIGN_OR_RETURN( - int64_t split_dim, - bitcast_split_fn( - replacement_instruction, - mutable_split_shape()->layout().split_configs(0).dimension())); - SplitConfig split_config( - split_dim, - {mutable_split_shape()->layout().split_configs(0).split_indices( - 0)}); - replacement_instruction->mutable_shape() - ->mutable_layout() - ->add_split_configs(split_config); - } - } - ABSL_RETURN_IF_ERROR(use.instruction->ReplaceOperandWith( - use.operand_number, replacement_instruction)); - } - return absl::OkStatus(); -} - -bool Allocation::is_copy_like_allocation() const { - return is_copy_allocation() || is_sliced_copy_allocation(); -} - -HloInstruction* Allocation::AddGetTupleElements() const { - CHECK_NE(defining_position().instruction, nullptr); - - Shape shape = defining_position().shape(); - CHECK(shape.IsArray()) << "Allocation shape is not an array. Shape = " - << shape.ToString() - << " position = " << defining_position().shape(); - return TupleUtil::AddGetTupleElements(defining_position()); -} - -Allocation::Allocation(HloPosition defining_position, MemorySpace memory_space, - std::optional chunk, - int64_t start_time, int64_t end_time, - std::optional cross_program_prefetch_index) - : original_defining_position_(std::move(defining_position)), - memory_space_(memory_space), - chunk_(chunk), - start_time_(start_time), - end_time_(end_time), - cross_program_prefetch_index_(cross_program_prefetch_index), - split_shape_(std::nullopt) {} - -HloPosition Allocation::original_defining_position() const { - return original_defining_position_; -} - -void Allocation::set_original_defining_position(HloPosition defining_position) { - original_defining_position_ = std::move(defining_position); -} - -bool Allocation::base_is_equal(const Allocation& other) const { - return defining_position() == other.defining_position() && - uses() == other.uses() && memory_space() == other.memory_space() && - chunk() == other.chunk() && start_time() == other.start_time() && - end_time() == other.end_time() && - earliest_available_time() == other.earliest_available_time(); -} - -PinnedAllocation::PinnedAllocation(HloPosition defining_position, - MemorySpace memory_space, - std::optional chunk, - int64_t start_time, int64_t end_time) - : Allocation(std::move(defining_position), memory_space, chunk, start_time, - end_time, - /*cross_program_prefetch_index=*/std::nullopt) {} - -HloPosition PinnedAllocation::defining_position() const { - return original_defining_position(); -} - -bool PinnedAllocation::operator==(const PinnedAllocation& other) const { - return this->base_is_equal(static_cast(other)); -} - -bool MirroredAllocation::operator==(const MirroredAllocation& other) const { - return this->base_is_equal(static_cast(other)); -} - -bool ParentAllocation::operator==(const ParentAllocation& other) const { - return this->base_is_equal(static_cast(other)); -} - -bool PinnedAllocation::operator==(const Allocation& other) const { - const PinnedAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -absl::Status PinnedAllocation::Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - HloInstruction* producing_instruction = AddGetTupleElements(); - HloComputation* computation = producing_instruction->parent(); - return UpdateUses(computation, producing_instruction, bitcast_split_fn, - hlo_live_range, alias_analysis); -} - -std::string PinnedAllocation::ToString() const { - std::string memory_space_str = MemorySpaceToString(memory_space()); - std::optional chunk = maybe_chunk(); - if (chunk) { - absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, - ", size: ", chunk->size, ")"); - } - return absl::StrCat("PinnedAllocation in ", memory_space_str, " defined at ", - original_defining_position().ToString(), - ", start_time:", start_time(), ", end_time:", end_time(), - ", uses: ", UsesToString(uses())); -} - -void PinnedAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void PinnedAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - needed_allocations.insert(this); -} - -ReservedAllocation::ReservedAllocation(HloPosition defining_position, - HeapSimulator::Chunk chunk, - int64_t start_time, int64_t end_time) - : Allocation(std::move(defining_position), MemorySpace::kAlternate, chunk, - /*start_time=*/start_time, /*end_time=*/end_time, - /*cross_program_prefetch_index=*/std::nullopt), - reserved_(true) {} - -HloPosition ReservedAllocation::defining_position() const { - return original_defining_position(); -} - -absl::Status ReservedAllocation::Process( - const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - return absl::OkStatus(); -} - -void ReservedAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void ReservedAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - needed_allocations.insert(this); -} - -std::string ReservedAllocation::ToString() const { - std::string memory_space_str = MemorySpaceToString(memory_space()); - std::optional chunk = maybe_chunk(); - if (chunk) { - absl::StrAppend(&memory_space_str, " chunk: ", chunk->ToString()); - } - return absl::StrCat("ReservedAllocation defined in ", memory_space_str, - " defined at ", original_defining_position().ToString(), - ", start_time: ", start_time(), - ", end_time: ", end_time(), " reserved: ", reserved_); -} - -bool ReservedAllocation::operator==(const Allocation& other) const { - const ReservedAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -bool ReservedAllocation::operator==(const ReservedAllocation& other) const { - return this->base_is_equal(static_cast(other)) && - reserved_ == other.reserved_; -} - -CopyAllocation::CopyAllocation( - Allocation& prev_allocation, MemorySpace memory_space, - std::optional chunk, - int64_t copy_start_schedule_after_time, - int64_t copy_done_schedule_before_time, int64_t end_time, - std::optional cross_program_prefetch_index, - HloInstruction* sync_mem_op, HloInstruction* async_mem_op_start, - HloInstruction* async_mem_op_done, int64_t source_operand_index) - : Allocation( - /*defining_position=*/{nullptr, {}}, memory_space, chunk, - // Allocation uses an inclusive start time - ExclusiveToInclusiveStartTime(copy_start_schedule_after_time), - end_time, cross_program_prefetch_index), - prev_allocation_(prev_allocation), - copy_start_schedule_after_(copy_start_schedule_after_time), - copy_done_schedule_before_(copy_done_schedule_before_time), - copy_start_(async_mem_op_start), - copy_done_(async_mem_op_done), - sync_mem_op_(sync_mem_op), - source_operand_index_(source_operand_index) {} - -int64_t CopyAllocation::earliest_available_time() const { - return copy_done_schedule_before_; -} - -absl::Status CopyAllocation::Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - // If the copy start and copy done instructions are already present in the - // original schedule we just need to schedule them in correct positions. There - // is no need to create new instructions. - if (copy_start_ != nullptr && copy_done_ != nullptr) { - // Update the allocation position with the copy complete instruction, so - // that if there are further copies from it, they can find the correct - // position. - set_original_defining_position(HloPosition{copy_start_, {0}}); - return absl::OkStatus(); - } - - // Copy allocations need to insert asynchronous copy nodes. - Shape shape = defining_position().shape(); - HloInstruction* producing_instruction = AddGetTupleElements(); - HloComputation* computation = producing_instruction->parent(); - VLOG(3) << "Processing copy allocation: " << ToString(); - if (sync_mem_op_ != nullptr && sync_mem_op_->opcode() != HloOpcode::kCopy) { - if (sync_mem_op_->opcode() == HloOpcode::kSlice || - sync_mem_op_->opcode() == HloOpcode::kDynamicSlice || - sync_mem_op_->IsCustomFusion()) { - ABSL_ASSIGN_OR_RETURN(copy_done_, - computation->CreateAsyncInstructions( - sync_mem_op_, {ShapeUtil::MakeShape(S32, {})}, - HloInstruction::kMainExecutionThread, false)); - } else { - return Internal("Sync mem op is not a copy, slice, or dynamic slice."); - } - copy_start_ = copy_done_->mutable_operand(0); - // If the shape of the copy start operand is not compatible with the - // shape of the producing instruction, we insert a bitcast to make them - // compatible. - if (!ShapeUtil::CompatibleIgnoringFpPrecision( - producing_instruction->shape(), - copy_start_->operand(source_operand_index_)->shape())) { - producing_instruction = - computation->AddInstruction(HloInstruction::CreateBitcast( - copy_start_->operand(source_operand_index_)->shape(), - producing_instruction)); - } - ABSL_RETURN_IF_ERROR(copy_start_->ReplaceOperandWith(source_operand_index_, - producing_instruction)); - } else { - Shape dest_shape = shape; - if (memory_space() == MemorySpace::kDefault) { - dest_shape.mutable_layout()->clear_split_configs(); - } - copy_start_ = computation->AddInstruction(HloInstruction::CreateCopyStart( - ShapeUtil::MakeTupleShape( - {dest_shape, shape, ShapeUtil::MakeShape(U32, {})}), - producing_instruction, cross_program_prefetch_index())); - copy_done_ = computation->AddInstruction(HloInstruction::CreateUnary( - dest_shape, HloOpcode::kCopyDone, copy_start_)); - } - VLOG(4) << "Created " << copy_start_->name() - << " for copy allocation: " << ToString(); - - // Update the allocation position with the copy complete instruction, so that - // if there are further copies from it, they can find the correct position. - set_original_defining_position(HloPosition{copy_done_, {}}); - return UpdateUses(computation, copy_done_, bitcast_split_fn, hlo_live_range, - alias_analysis); -} - -void CopyAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void CopyAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - if (needed_allocations.insert(this).second) { - prev_allocation_.MarkNeeded(needed_allocations); - } -} - -std::string CopyAllocation::ToString() const { - std::string memory_space_str = MemorySpaceToString(memory_space()); - std::optional chunk = maybe_chunk(); - if (chunk) { - absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, - ", size: ", chunk->size, ")"); - } - return absl::StrCat( - "Copy Allocation in ", memory_space_str, ", start_time:", start_time(), - ", end_time:", end_time(), - ", copy_start_after_time: ", copy_start_schedule_after(), - ", copy_done_before_time: ", copy_done_schedule_before(), - ", uses: ", UsesToString(uses()), - ", existing sync_mem_op: ", sync_mem_op_ ? sync_mem_op_->name() : "none", - ", existing copy_start: ", copy_start_ ? copy_start_->name() : "none", - ", existing copy_done: ", copy_done_ ? copy_done_->name() : "none", - ", from ", prev_allocation_.ToString()); -} - -HloPosition CopyAllocation::defining_position() const { - // Unless explicitly set, the defining position of a copy allocation is - // retrieved from the previous allocation. This is because we don't create - // new CopyStart/CopyDone instructions until later and the position should - // point to the previous (copy or otherwise) allocation's position for the - // original defining position. - HloPosition defining_position = original_defining_position(); - if (defining_position.instruction == nullptr) { - return prev_allocation_.defining_position(); - } - return defining_position; -} - -bool CopyAllocation::operator==(const CopyAllocation& other) const { - return this->base_is_equal(static_cast(other)) && - copy_done_schedule_before() == other.copy_done_schedule_before() && - copy_start_schedule_after() == other.copy_start_schedule_after() && - copy_start() == other.copy_start() && copy_done() == other.copy_done(); -} - -bool CopyAllocation::operator==(const Allocation& other) const { - const CopyAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -void CopyAllocation::set_copy_start_schedule_after( - int64_t copy_start_schedule_after) { - copy_start_schedule_after_ = copy_start_schedule_after; -} - -void CopyAllocation::set_copy_done_schedule_before( - int64_t copy_done_schedule_before) { - copy_done_schedule_before_ = copy_done_schedule_before; -} - -int64_t CopyAllocation::copy_start_schedule_after() const { - return copy_start_schedule_after_; -} - -int64_t CopyAllocation::copy_done_schedule_before() const { - return copy_done_schedule_before_; -} - -SlicedCopyAllocation::SlicedCopyAllocation( - const Allocation& prev_allocation, MemorySpace memory_space, - std::vector slice_decisions_sorted_by_exclusive_start_time, - int64_t copy_done_schedule_before_time, int64_t end_time, - const SlicedPrefetchOptions& sliced_prefetch_options, - absl::FunctionRef get_equivalent_s8_shape_fn, - HloInstruction* sync_mem_op) - : Allocation( - /*defining_position=*/{nullptr, {}}, memory_space, - GetSlicedCopyAllocationChunk( - slice_decisions_sorted_by_exclusive_start_time), - // Allocation uses an inclusive start time - ExclusiveToInclusiveStartTime( - GetSlicedCopyAllocationExclusiveStartTime( - slice_decisions_sorted_by_exclusive_start_time)), - end_time, - /*cross_program_prefetch_index=*/std::nullopt), - original_shape_to_slice_(prev_allocation.defining_position().shape()), - prev_allocation_(prev_allocation), - sliced_prefetch_options_(sliced_prefetch_options), - get_equivalent_s8_shape_fn_(get_equivalent_s8_shape_fn), - sync_mem_op_(sync_mem_op) { - CHECK_GE(slice_decisions_sorted_by_exclusive_start_time.size(), 2); - slice_details_sorted_by_exclusive_start_time_.reserve( - slice_decisions_sorted_by_exclusive_start_time.size()); - for (SliceDecision& decision : - slice_decisions_sorted_by_exclusive_start_time) { - int64_t copy_done_schedule_after_time = decision.exclusive_start_time; - slice_details_sorted_by_exclusive_start_time_.push_back(SliceDetail{ - std::move(decision), - copy_done_schedule_after_time, - copy_done_schedule_before_time, - /*copy_start=*/nullptr, - /*copy_done=*/nullptr, - }); - } -} - -absl::Status SlicedCopyAllocation::Process( - const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - Shape shape = defining_position().shape(); - HloInstruction* producing_instruction = AddGetTupleElements(); - - // Calling Process() over the previous allocation might - // have modified the defining position, and hence the shape that was used when - // we computed the slices. In cases where the shape has changed, we insert a - // bitcast, so slice instructions operate on the originally sliced shape. - // - // Note, these bitcasts are being inserted in the same cases that - // UpdateUses() is inserting bitcasts, except we are - // inserting the bitcasts before the copy, instead of after the copy. - if (!Shape::Equal().IgnoreMemorySpaceInLayout()(shape, - original_shape_to_slice_)) { - int64_t new_memory_space = shape.layout().memory_space(); - shape = original_shape_to_slice_; - shape.mutable_layout()->set_memory_space(new_memory_space); - producing_instruction = producing_instruction->parent()->AddInstruction( - HloInstruction::CreateBitcast(shape, producing_instruction)); - } - - HloComputation* computation = producing_instruction->parent(); - std::vector slice_dones; - slice_dones.reserve(slice_details_sorted_by_exclusive_start_time_.size()); - - // If we are trying to make all slices a uniform size, we bitcast the - // producing instruction to an array of bytes, so it is easy to slice into any - // size. - Shape slice_shape = shape; - if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) { - slice_shape = get_equivalent_s8_shape_fn_(shape); - producing_instruction = producing_instruction->parent()->AddInstruction( - HloInstruction::CreateBitcast(slice_shape, producing_instruction)); - } - - // Sliced copy allocations need to insert asynchronous copy nodes. - for (SliceDetail& slice_detail : - slice_details_sorted_by_exclusive_start_time_) { - ABSL_RETURN_IF_ERROR(slice_detail.CreateAsyncSlice( - slice_shape, *producing_instruction, *computation)); - VLOG(4) << "Created " << slice_detail.copy_start->name() - << " for sliced copy allocation: " << ToString(); - slice_dones.push_back(slice_detail.copy_done); - } - - ABSL_RETURN_IF_ERROR(CreateBitcastConcat(shape, slice_dones)); - - // 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 - // concatenation back to the original shape. - if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) { - concat_ = concat_->parent()->AddInstruction( - HloInstruction::CreateBitcast(shape, concat_)); - } - - // Update the allocation position with the copy complete instruction, so that - // if there are further copies from it, they can find the correct position. - set_original_defining_position(HloPosition{concat_, {}}); - return UpdateUses(computation, concat_, bitcast_split_fn, hlo_live_range, - alias_analysis); -} - -void SlicedCopyAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void SlicedCopyAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - if (needed_allocations.insert(this).second) { - prev_allocation_.MarkNeeded(needed_allocations); - } -} - -HloPosition SlicedCopyAllocation::defining_position() const { - // Unless explicitly set, the defining position of a sliced copy allocation is - // retrieved from the previous allocation. This is because we don't create - // new CopyStart/CopyDone instructions until later and the position should - // point to the previous (copy or otherwise) allocation's position for the - // original defining position. - HloPosition defining_position = original_defining_position(); - if (defining_position.instruction == nullptr) { - return prev_allocation_.defining_position(); - } - return defining_position; -} - -int64_t SlicedCopyAllocation::earliest_available_time() const { - return slice_details_sorted_by_start_time().back().copy_done_before_time; -} - -std::vector SlicedCopyAllocation::SliceOffsetsSortedByStartTime() - const { - std::vector offsets; - offsets.reserve(slice_details_sorted_by_exclusive_start_time_.size()); - - for (const SliceDetail& slice_detail : - slice_details_sorted_by_exclusive_start_time_) { - offsets.push_back(slice_detail.slice_decision.chunk.offset); - } - - return offsets; -} - -void SlicedCopyAllocation::AddDiffToAllSliceOffsets(int64_t diff) { - for (SliceDetail& slice_detail : - slice_details_sorted_by_exclusive_start_time_) { - HeapSimulator::Chunk& chunk = slice_detail.slice_decision.chunk; - chunk = - HeapSimulator::Chunk::FromOffsetSize(chunk.offset + diff, chunk.size); - } -} - -void SlicedCopyAllocation::ImportRepackedSliceData( - const SlicedAllocationData& data) { - int num_slices = slice_details_sorted_by_exclusive_start_time_.size(); - CHECK_EQ(data.slices_sorted_by_offset.size(), num_slices); - - std::vector slice_details_sorted_by_offset; - slice_details_sorted_by_offset.reserve(num_slices); - for (SliceDetail& slice_detail : - slice_details_sorted_by_exclusive_start_time_) { - slice_details_sorted_by_offset.push_back(&slice_detail); - } - absl::c_sort(slice_details_sorted_by_offset, [](const SliceDetail* lhs, - const SliceDetail* rhs) { - return lhs->slice_decision.chunk.offset < rhs->slice_decision.chunk.offset; - }); - - for (int i = 0; i < num_slices; ++i) { - SliceDetail* slice_detail = slice_details_sorted_by_offset[i]; - HeapSimulator::Chunk& chunk = slice_detail->slice_decision.chunk; - const AllocatedSlice& repacked_slice_data = data.slices_sorted_by_offset[i]; - chunk = HeapSimulator::Chunk::FromOffsetSize(repacked_slice_data.offset, - chunk.size); - slice_detail->copy_start_after_time = - repacked_slice_data.inclusive_start_time - 1; - slice_detail->slice_decision.exclusive_start_time = - InclusiveToExclusiveStartTime(repacked_slice_data.inclusive_start_time); - } - - absl::c_sort(slice_details_sorted_by_exclusive_start_time_, - [](const SliceDetail& lhs, const SliceDetail& rhs) { - return std::make_tuple(lhs.copy_start_after_time, - lhs.slice_decision.chunk.offset) < - std::make_tuple(rhs.copy_start_after_time, - rhs.slice_decision.chunk.offset); - }); -} - -const std::vector& -SlicedCopyAllocation::slice_details_sorted_by_start_time() const { - return slice_details_sorted_by_exclusive_start_time_; -} - -std::vector& -SlicedCopyAllocation::mutable_slice_details_sorted_by_start_time() { - return slice_details_sorted_by_exclusive_start_time_; -} - -bool SlicedCopyAllocation::operator==(const SlicedCopyAllocation& other) const { - return this->base_is_equal(static_cast(other)) && - slice_details_sorted_by_exclusive_start_time_ == - other.slice_details_sorted_by_exclusive_start_time_ && - concat_ == other.concat_; -} - -std::string SlicedCopyAllocation::ToString() const { - std::string memory_space_str = MemorySpaceToString(memory_space()); - if (memory_space() == MemorySpace::kAlternate) { - absl::StrAppend(&memory_space_str, " (off: ", maybe_chunk()->offset, ")"); - } - return absl::StrCat( - "Sliced Copy Allocation in ", memory_space_str, - ", start_time:", start_time(), ", end_time:", end_time(), - ", first_slice_copy_start_after_time: ", - slice_details_sorted_by_start_time().front().copy_start_after_time, - ", last_slice_copy_done_before_time: ", - slice_details_sorted_by_start_time().back().copy_done_before_time, - ", uses: ", UsesToString(uses()), - ", sync_mem_op: ", sync_mem_op_ ? sync_mem_op_->name() : "none", - ", from ", prev_allocation_.ToString()); -} - -absl::Status SlicedCopyAllocation::CreateBitcastConcat( - const Shape& shape, absl::Span slices) { - CHECK(!slices.empty()); - concat_ = - slices.front()->parent()->AddInstruction(HloInstruction::CreateCustomCall( - shape, slices, - xla::memory_space_assignment::kConcatBitcastCustomCall)); - return absl::OkStatus(); -} - -std::string SlicedCopyAllocation::SliceDetail::ToString() const { - return absl::StrCat("{ slice_decision: ", slice_decision.ToString(), - ", copy_start_after_time: ", copy_start_after_time, - ", copy_done_before_time: ", copy_done_before_time, " }"); -} - -std::tuple -SliceDetailToTuple(const SlicedCopyAllocation::SliceDetail& slice_detail) { - return std::make_tuple(std::ref(slice_detail.slice_decision), - slice_detail.copy_start_after_time, - slice_detail.copy_done_before_time, - slice_detail.copy_start, slice_detail.copy_done); -} - -bool SlicedCopyAllocation::SliceDetail::operator==( - const SliceDetail& other) const { - return SliceDetailToTuple(*this) == SliceDetailToTuple(other); -} - -absl::Status SlicedCopyAllocation::SliceDetail::CreateAsyncSlice( - const Shape& original_shape, HloInstruction& producer, - HloComputation& parent) { - if (original_shape.dimensions().size() != - slice_decision.sizing.slice_params.size()) { - return FailedPrecondition( - "%s", absl::StrCat("The number of SlicedCopyAllocation parameters ", - slice_decision.sizing.slice_params.size(), - " does not match the rank ", - original_shape.dimensions().size(), - " of the tensor we are slicing.")); - } - - std::vector start_indices; - start_indices.reserve(slice_decision.sizing.slice_params.size()); - std::vector limit_indices; - limit_indices.reserve(slice_decision.sizing.slice_params.size()); - std::vector strides; - strides.reserve(slice_decision.sizing.slice_params.size()); - - for (int i = 0; i < slice_decision.sizing.slice_params.size(); ++i) { - const SliceParam& slice_param = slice_decision.sizing.slice_params[i]; - start_indices.push_back(slice_param.start_inclusive); - limit_indices.push_back(slice_param.end_exclusive); - strides.push_back(1); - const int64_t new_dim = - slice_param.end_exclusive - slice_param.start_inclusive; - if (new_dim <= 0) { - return FailedPrecondition( - "%s", absl::StrCat("SlicedCopyAllocation new dimension size is ", - new_dim, ", expected something > 0.")); - } - if (original_shape.dimensions(i) < new_dim) { - return FailedPrecondition( - "%s", - absl::StrCat("SlicedCopyAllocation sliced dimension size ", new_dim, - " is bigger than its original dimension size of ", - original_shape.dimensions(i), ".")); - } - } - - HloInstruction* slice = parent.AddInstruction( - HloInstruction::CreateSlice(slice_decision.sizing.slice_shape, &producer, - start_indices, limit_indices, strides)); - ABSL_ASSIGN_OR_RETURN(copy_done, parent.CreateAsyncInstructions( - slice, {ShapeUtil::MakeShape(S32, {})})); - copy_start = copy_done->mutable_operand(0); - - return absl::OkStatus(); -} - -bool SlicedCopyAllocation::operator==(const Allocation& other) const { - const SlicedCopyAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -MirroredAllocation::MirroredAllocation(const Allocation& original_allocation, - int64_t time) - : Allocation(original_allocation.defining_position(), - original_allocation.memory_space(), - original_allocation.maybe_chunk(), - /*start_time=*/time, - /*end_time=*/time, - /*cross_program_prefetch_index=*/std::nullopt), - defining_position_(std::nullopt), - original_allocation_(original_allocation) {} - -MirroredAllocation::MirroredAllocation(HloPosition defining_position, - const Allocation& original_allocation, - int64_t start_time, int64_t end_time) - : Allocation(original_allocation.defining_position(), - original_allocation.memory_space(), - original_allocation.maybe_chunk(), - /*start_time=*/start_time, - /*end_time=*/end_time, - /*cross_program_prefetch_index=*/std::nullopt), - defining_position_(defining_position), - original_allocation_(original_allocation) {} - -HloPosition MirroredAllocation::defining_position() const { - if (defining_position_.has_value()) { - return defining_position_.value(); - } - return original_allocation_.defining_position(); -} - -std::string MirroredAllocation::ToString() const { - std::string memory_space_str = MemorySpaceToString(memory_space()); - std::optional chunk = maybe_chunk(); - if (chunk) { - absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, - ", size: ", chunk->size, ")"); - } - return absl::StrCat( - "MirroredAllocation in ", memory_space_str, " defined at ", - defining_position().ToString(), ", start_time:", start_time(), - ", end_time:", end_time(), ", uses: ", UsesToString(uses()), - ", for original allocation: ", original_allocation_.ToString()); -} - -absl::Status MirroredAllocation::Process( - const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - set_original_defining_position(defining_position()); - HloInstruction* producing_instruction = AddGetTupleElements(); - HloComputation* computation = producing_instruction->parent(); - return UpdateUses(computation, producing_instruction, bitcast_split_fn, - hlo_live_range, alias_analysis); -} - -void MirroredAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void MirroredAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - if (needed_allocations.insert(this).second) { - original_allocation_.MarkNeeded(needed_allocations); - } -} - -bool MirroredAllocation::operator==(const Allocation& other) const { - const MirroredAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other) && - defining_position() == other.defining_position(); -} - -ParentAllocation::ParentAllocation(const Allocation& original_allocation, - HloInstruction* calling_instruction, - HloPosition position, int64_t time) - : Allocation(std::move(position), MemorySpace::kDefault, - original_allocation.maybe_chunk(), - /*start_time=*/time, - /*end_time=*/time, - /*cross_program_prefetch_index=*/std::nullopt), - original_allocation_(original_allocation), - calling_instruction_(calling_instruction) {} - -HloPosition ParentAllocation::defining_position() const { - return original_defining_position(); -} - -std::string ParentAllocation::ToString() const { - return absl::StrCat("Parent Allocation mirrored at ", - original_defining_position().ToString(), ", originally ", - original_allocation_.ToString()); -} - -absl::Status ParentAllocation::Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - // Add an additional parameter to the while HLO with a reference to the buffer - // in the default memory space. - HloInstruction* producing_instruction = - original_allocation_.AddGetTupleElements(); - int new_tuple_index = calling_instruction_->shape().tuple_shapes().size(); - - ABSL_ASSIGN_OR_RETURN( - HloInstruction * new_while_operand, - TupleUtil::ReplaceTupleWith(producing_instruction, - calling_instruction_->mutable_operand(0), - {new_tuple_index})); - ABSL_RETURN_IF_ERROR(calling_instruction_->ReplaceOperandWithDifferentShape( - 0, new_while_operand)); - *calling_instruction_->mutable_shape() = new_while_operand->shape(); - *calling_instruction_->while_condition() - ->parameter_instruction(0) - ->mutable_shape() = new_while_operand->shape(); - *calling_instruction_->while_body() - ->parameter_instruction(0) - ->mutable_shape() = new_while_operand->shape(); - HloPosition defining_position = original_defining_position(); - defining_position.index = {new_tuple_index}; - set_original_defining_position(defining_position); - // Also replace the while op with a tuple that has the old shape. Note that we - // need to first take a snapshot of the users before calling ExtractPrefix - // since ExtractPrefix introduces additional gte users. - std::vector while_users = calling_instruction_->users(); - HloInstruction* tuple_with_old_shape = - TupleUtil::ExtractPrefix(calling_instruction_, new_tuple_index); - ABSL_RETURN_IF_ERROR(calling_instruction_->ReplaceAllUsesWithDifferentShape( - while_users, tuple_with_old_shape)); - - HloInstruction* final_instruction = AddGetTupleElements(); - HloComputation* computation = final_instruction->parent(); - return UpdateUses(computation, final_instruction, bitcast_split_fn, - hlo_live_range, alias_analysis); -} - -absl::Status ParentAllocation::PostProcess() { - // Update the root of the while body with the new parameter. The reason why we - // need a separate post-process for this is because other allocations may have - // while body root as a use, so they would update the old root instead of the - // new root. Doing the post-process step later ensures the root has been - // updated with other changes, and we can safely add the additional parameter. - HloComputation* while_body = calling_instruction_->while_body(); - ABSL_ASSIGN_OR_RETURN(HloInstruction * new_while_body_root, - TupleUtil::ReplaceTupleWith( - AddGetTupleElements(), while_body->root_instruction(), - original_defining_position().index)); - while_body->set_root_instruction(new_while_body_root, - /*accept_different_shape=*/true); - return absl::OkStatus(); -} - -void ParentAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - // Parent allocations are only needed if they have any uses or if there is a - // copy allocation that copies this value (in that case, the copy allocation - // will call this allocation's MarkNeeded function). - if (!has_no_uses()) { - MarkNeeded(needed_allocations); - } -} - -void ParentAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - if (needed_allocations.insert(this).second) { - original_allocation_.MarkNeeded(needed_allocations); - } -} - -bool ParentAllocation::operator==(const Allocation& other) const { - const ParentAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -WindowPrefetchedAllocation::WindowPrefetchedAllocation( - Allocation& prev_allocation, HloUse use, const HeapSimulator::Chunk& chunk, - int64_t prefetch_start_schedule_after_time, - int64_t prefetch_done_schedule_before_time, const Options& options) - : Allocation( - {nullptr, {}}, MemorySpace::kAlternate, chunk, - ExclusiveToInclusiveStartTime(prefetch_start_schedule_after_time), - InclusiveToExclusiveEndTime(prefetch_done_schedule_before_time), - /*cross_program_prefetch_index=*/std::nullopt), - options_(options), - defining_position_(prev_allocation.defining_position()), - use_(use), - prefetch_start_schedule_after_(prefetch_start_schedule_after_time), - prefetch_done_schedule_before_(prefetch_done_schedule_before_time), - bytes_(chunk.size) {} - -HloPosition WindowPrefetchedAllocation::defining_position() const { - if (defining_position_.instruction != nullptr) { - return defining_position_; - } - return original_defining_position(); -} - -int64_t WindowPrefetchedAllocation::earliest_available_time() const { - return prefetch_done_schedule_before_; -} - -absl::Status WindowPrefetchedAllocation::InsertWindowPrefetchInstruction( - HloInstruction* producing_instruction, HloInstruction* use_instruction, - HloComputation* computation) { - // Derive the shape for window buffer. - Shape buffer_shape = ShapeUtil::MakeShape(U8, {options_.bytes}); - Layout layout = LayoutUtil::MakeLayout({0}); - layout.set_memory_space(options_.alternate_memory_space); - *buffer_shape.mutable_layout() = layout; - // Sync flag shape - Shape sflag_shape = ShapeUtil::MakeShape(S32, {2}); - // Output shape of the WindowPrefetch op. - Shape output_shape = ShapeUtil::MakeTupleShape({buffer_shape, sflag_shape}); - - // Insert WindowPrefetch op. - HloInstruction* custom_call = - computation->AddInstruction(HloInstruction::CreateCustomCall( - output_shape, {producing_instruction}, "WindowPrefetch")); - HloInstruction* get_buffer = computation->AddInstruction( - HloInstruction::CreateGetTupleElement(buffer_shape, custom_call, 0)); - HloInstruction* get_sflag = computation->AddInstruction( - HloInstruction::CreateGetTupleElement(sflag_shape, custom_call, 1)); - use_instruction->AppendOperand(get_buffer); - use_instruction->AppendOperand(get_sflag); - - // Add backend config for the WindowPrefetch op. - int64_t parameter_index = - use_instruction->operand_index(producing_instruction); - int64_t buffer_index = use_instruction->operand_index(get_buffer); - custom_call->set_raw_backend_config_string(absl::StrCat( - "parameter_index=", parameter_index, " buffer_index=", buffer_index)); - - // The buffer's defining position is the get_tuple_element instruction. - prefetch_instruction_ = get_buffer; - return absl::OkStatus(); -} - -absl::Status WindowPrefetchedAllocation::Process( - const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - HloInstruction* use_instruction = use_.instruction; - int64_t use_operand = use_instruction->operand_count(); - HloInstruction* producing_instruction = - use_instruction->mutable_operand(use_.operand_number); - HloComputation* computation = producing_instruction->parent(); - CHECK_EQ(use_instruction->opcode(), HloOpcode::kFusion); - - ABSL_RETURN_IF_ERROR(InsertWindowPrefetchInstruction( - producing_instruction, use_instruction, computation)); - - // Notify the backend that an operand has been appended as a window prefetch - // buffer. - options_.notify_operand_appended_fn(use_instruction, use_.operand_number, - use_operand); - - // Set the original defining position to the window prefetch instruction. - set_original_defining_position(defining_position_); - defining_position_ = {prefetch_instruction_, {}}; - AddUse(HloUse{use_instruction, use_operand}); - return absl::OkStatus(); -} - -void WindowPrefetchedAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void WindowPrefetchedAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - needed_allocations.insert(this); -} - -std::string WindowPrefetchedAllocation::ToString() const { - return absl::StrCat("WindowPrefetched Allocation"); -} - -bool WindowPrefetchedAllocation::operator==( - const WindowPrefetchedAllocation& other) const { - return this->base_is_equal(static_cast(other)) && - prefetch_done_schedule_before() == - other.prefetch_done_schedule_before() && - prefetch_start_schedule_after() == - other.prefetch_start_schedule_after() && - prefetch() == other.prefetch() && bytes_ == other.bytes_; -} - -bool WindowPrefetchedAllocation::operator==(const Allocation& other) const { - const WindowPrefetchedAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -ScopedAllocation::ScopedAllocation(HeapSimulator::Chunk chunk, - int64_t allocation_time, - HloInstruction* defining_instruction, - bool is_post_module) - : Allocation(/*defining_position=*/{defining_instruction, {}}, - MemorySpace::kAlternate, chunk, - /*start_time=*/allocation_time, - /*end_time=*/allocation_time, - /*cross_program_prefetch_index=*/std::nullopt), - is_post_module_(is_post_module) {} - -HloPosition ScopedAllocation::defining_position() const { - return original_defining_position(); -} - -absl::Status ScopedAllocation::Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - return absl::OkStatus(); -} - -void ScopedAllocation::MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const { - MarkNeeded(needed_allocations); -} - -void ScopedAllocation::MarkNeeded( - absl::flat_hash_set& needed_allocations) const { - needed_allocations.insert(this); -} - -std::string ScopedAllocation::ToString() const { - std::string name = ""; - if (!is_post_module_) { - name = defining_position().instruction->name(); - } - return absl::StrCat("Scoped Allocation defined for ", name, - ", allocation_time: ", start_time(), - ", chunk: ", chunk().ToString()); -} - -bool ScopedAllocation::operator==(const Allocation& other) const { - const ScopedAllocation* casted_other = - dynamic_cast(&other); - return casted_other != nullptr && (*this) == (*casted_other); -} - -bool ScopedAllocation::operator==(const ScopedAllocation& other) const { - return this->base_is_equal(static_cast(other)) && - is_post_module_ == other.is_post_module_; -} - -std::tuple GetAllocationSortTuple( - const std::unique_ptr& allocation) { - int64_t scheduled_on_or_before = allocation->start_time(); - int64_t scheduled_on_or_after = allocation->start_time(); - if (allocation->is_copy_allocation()) { - auto copy_allocation = absl::down_cast(allocation.get()); - scheduled_on_or_before = copy_allocation->copy_done_schedule_before(); - scheduled_on_or_after = copy_allocation->copy_start_schedule_after(); - } - return std::forward_as_tuple(scheduled_on_or_before, - !allocation->is_copy_allocation(), - scheduled_on_or_after); -} - -void SortAllocationSequence(AllocationSequence& allocations) { - absl::c_sort(allocations, [](const std::unique_ptr& lhs, - const std::unique_ptr& rhs) { - return GetAllocationSortTuple(lhs) < GetAllocationSortTuple(rhs); - }); -} - -std::string AllocationSequenceToString(AllocationSequence& allocations, - bool sort_allocations) { - if (sort_allocations) { - SortAllocationSequence(allocations); - } - std::string allocations_str = "\n"; - for (const std::unique_ptr& allocation : allocations) { - absl::StrAppend(&allocations_str, allocation->ToString(), "\n"); - } - return allocations_str; -} - -std::vector GetAllocationSequenceInRawPointers( - AllocationSequence& allocations) { - std::vector allocations_in_raw_pointers; - for (const std::unique_ptr& allocation : allocations) { - allocations_in_raw_pointers.push_back(allocation.get()); - } - return allocations_in_raw_pointers; -} - -namespace { - -struct AllocationSummary { - static void Add(const Allocation& allocation, - std::vector& data) { - if (!allocation.is_sliced_copy_allocation()) { - data.push_back( - AllocationSummary{allocation.chunk(), allocation.start_time(), - allocation.end_time(), GetName(allocation)}); - return; - } - const SlicedCopyAllocation& sliced_copy_allocation = - dynamic_cast(allocation); - for (size_t i = 0; - i < sliced_copy_allocation.slice_details_sorted_by_start_time().size(); - ++i) { - const SlicedCopyAllocation::SliceDetail& slice_detail = - sliced_copy_allocation.slice_details_sorted_by_start_time()[i]; - data.push_back(AllocationSummary{ - slice_detail.slice_decision.chunk, - ExclusiveToInclusiveStartTime( - slice_detail.slice_decision.exclusive_start_time), - sliced_copy_allocation.end_time(), GetName(allocation, i)}); - } - } - - static std::string GetName(const Allocation& allocation, - std::optional slice_index = std::nullopt) { - std::vector attributes; - if (slice_index.has_value()) { - attributes.push_back(absl::StrCat("slice ", *slice_index)); - } - if (allocation.is_scoped_allocation()) { - attributes.push_back("scoped"); - } - if (allocation.cross_program_prefetch_index().has_value()) { - attributes.push_back("cross-program prefetch"); - } - if (attributes.empty()) { - return allocation.defining_position().ToString(); - } - return absl::StrCat(allocation.defining_position().ToString(), " (", - absl::StrJoin(attributes, ", "), ")"); - } - - HeapSimulator::Chunk chunk; - int64_t start_time_inclusive; - int64_t end_time_inclusive; - std::string name; -}; - -} // namespace - -void AllocationSequenceDebugging::LogAltMemAllocationsAt( - const AllocationSequence& allocations, int64_t time) { - std::vector data_vector; - for (const std::unique_ptr& allocation : allocations) { - if (allocation->start_time() <= time && allocation->end_time() >= time && - allocation->is_in_alternate_mem()) { - AllocationSummary::Add(*allocation, data_vector); - } - } - absl::c_sort(data_vector, - [](const AllocationSummary& a, const AllocationSummary& b) { - if (a.chunk.offset != b.chunk.offset) { - return a.chunk.offset < b.chunk.offset; - } - return a.start_time_inclusive < b.start_time_inclusive; - }); - LOG(INFO) << "Live allocations in alternate mem at instruction time " << time - << " (before MSA alters the graph):"; - for (const AllocationSummary& data : data_vector) { - LOG(INFO) << "Alt mem allocation in chunk " << data.chunk.ToString() - << " during [" << data.start_time_inclusive << "," - << data.end_time_inclusive << "], holding " << data.name; - } -} - -} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.h b/third_party/xla/xla/service/memory_space_assignment/allocation.h deleted file mode 100644 index 1f9f8bc89250c1..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/allocation.h +++ /dev/null @@ -1,736 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALLOCATION_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALLOCATION_H_ - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_set.h" -#include "absl/functional/function_ref.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/shape.h" - -namespace xla::memory_space_assignment { - -// MemorySpaceAssignment uses a notion of a slow and large default memory -// space and a fast and small alternate memory space. -enum class MemorySpace : std::uint8_t { kDefault, kAlternate }; -std::string MemorySpaceToString(MemorySpace memory_space); - -using BitcastSplitFn = std::function( - const HloInstruction* instruction, int64_t split_dim)>; - -// An interface describing what to do with a value in memory over its lifetime. -// An allocation might either be placed in the default or alternate memory. An -// HloValue might live in multiple different allocations over its lifetime. The -// lifetimes of the allocations are defined using start_time and end_time, which -// corresponds to the instruction indexes in the flattened schedule. Each of -// these allocations might partially overlap with each other. -// -// Consider an instruction Foo, and its users Bar and Baz, and the times given -// in terms of the flattened schedule of the entire module: -// -// Foo:10 -// / \ -// Bar:14 \ -// Baz:25 -// -// A valid memory space assignment could be like the following: -// -// Time: 10 ... 14 ... 25 -// Foo Bar Baz -// Alternate +-------+ +-----+ -// Default +---------------------+ -// ^ ^ ^ ^ -// | | | | -// evict evict prefetch prefetch -// start end start end -// -// This would be represented with: -// - PinnedAllocation(memory_space=kAlternate, start_time=10, end_time=14) -// - CopyAllocation(memory_space=kDefault, start_time=12, end_time=25) -// - CopyAllocation(memory_space=kAlternate, start_time=22, end_time=25) -class Allocation { - public: - virtual ~Allocation() = default; - - // Allocation source methods - // -------------------------------------------------------------------------- - // Returns the defining position for this allocation. - virtual HloPosition defining_position() const = 0; - // Returns the cross-program prefetch index for this allocation. - std::optional cross_program_prefetch_index() const; - - void set_split_shape(const std::optional& split_shape) { - split_shape_ = split_shape; - } - const std::optional& split_shape() const { return split_shape_; } - std::optional& mutable_split_shape() { return split_shape_; } - - // Allocation timing methods - // -------------------------------------------------------------------------- - // TODO(cl/604356742): update all timing methods to explicitly state that - // they're representing inclusive intervals. - int64_t start_time() const { return start_time_; } - int64_t end_time() const { return end_time_; } - // Returns the time the buffer is first available to be used - virtual int64_t earliest_available_time() const = 0; - void set_start_time(int64_t start_time) { start_time_ = start_time; } - void set_end_time(int64_t end_time) { end_time_ = end_time; } - // Extends the end time of this allocation. - void Extend(int64_t end_time) { end_time_ = std::max(end_time_, end_time); } - - // Allocation space methods - // -------------------------------------------------------------------------- - MemorySpace memory_space() const { return memory_space_; } - // Returns the associated chunk that may be a nullopt if the allocation is - // in the default memory space. - virtual std::optional maybe_chunk() const { - return chunk_; - } - // Returns the associated chunk. The caller should ensure that the chunk is - // defined (the allocation should be in the alternate memory space). - virtual HeapSimulator::Chunk chunk() const; - HeapSimulator::Chunk* mutable_chunk() { return &*chunk_; } - void set_offset(int64_t offset); - // Returns true if the allocation is in the alternate memory space. - bool is_in_alternate_mem() const; - // Returns true if the allocation is in the default memory space. - bool is_in_default_mem() const; - - // Use methods - // -------------------------------------------------------------------------- - const std::vector& uses() const { return uses_; } - void clear_uses() { uses_.clear(); } - bool has_no_uses() const { return uses_.empty(); } - // Adds a use to this allocation. - void AddUse(HloUse use); - void RemoveUse(HloUse use); - // Replaces all uses of the allocation with the copy_complete instruction. - absl::Status UpdateUses(HloComputation* computation, - HloInstruction* producing_instruction, - const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis); - - // Allocation type methods - // -------------------------------------------------------------------------- - virtual bool is_pinned_allocation() const = 0; - virtual bool is_copy_allocation() const = 0; - virtual bool is_sliced_copy_allocation() const = 0; - virtual bool is_window_prefetched_allocation() const = 0; - virtual bool is_scoped_allocation() const = 0; - virtual bool is_reserved_allocation() const = 0; - virtual bool is_mirrored_allocation() const = 0; - // True if the allocation is for a copy or a sliced-copy. - bool is_copy_like_allocation() const; - - // Processing methods - // -------------------------------------------------------------------------- - // Recursively create kGetTupleElement instructions if the defining position - // shape is not an array. Returns the new instruction that has array shape. - HloInstruction* AddGetTupleElements() const; - // After all of the time ranges for the allocations have been assigned, - // Process morphs the instructions affected to assign the memory spaces and - // insert asynchronous copy instructions if necessary. - virtual absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) = 0; - // An optional post-process step that will be called after all allocations - // have been processed. - virtual absl::Status PostProcess() = 0; - // Marks (adds this allocation to needed_allocations) if this allocation is - // needed. PinnedAllocation and CopyAllocations are always needed and - // ParentAllocations are needed if they have any uses or if other - // CopyAllocation or ParentAllocations depend on them. - virtual void MarkIfNeeded( - absl::flat_hash_set& needed_allocations) const = 0; - // Marks this allocation as needed. - virtual void MarkNeeded( - absl::flat_hash_set& needed_allocations) const = 0; - - // Utility methods - // -------------------------------------------------------------------------- - virtual std::string ToString() const = 0; - virtual bool operator==(const Allocation& other) const = 0; - - protected: - // Protected constructor to encourage use of the final subclasses (e.g., - // PinnedAllocation, CopyAllocation, etc.). - Allocation(HloPosition defining_position, MemorySpace memory_space, - std::optional chunk, int64_t start_time, - int64_t end_time, - std::optional cross_program_prefetch_index); - - // Returns the original defining position of this allocation. - HloPosition original_defining_position() const; - // Sets the original defining position of this allocation. - void set_original_defining_position(HloPosition defining_position); - bool base_is_equal(const Allocation& other) const; - - private: - HloPosition original_defining_position_; - MemorySpace memory_space_; - std::optional chunk_; - int64_t start_time_; - int64_t end_time_; - std::vector uses_; - std::optional cross_program_prefetch_index_; - // If present, indicates the newly split shape. - std::optional split_shape_; -}; - -using AllocationSequence = std::vector>; -std::tuple GetAllocationSortTuple( - const std::unique_ptr& allocation); -void SortAllocationSequence(AllocationSequence& allocations); -std::string AllocationSequenceToString(AllocationSequence& allocations, - bool sort_allocations = false); -std::vector GetAllocationSequenceInRawPointers( - AllocationSequence& allocations); - -// This class represents an allocation that pins a tensor to -// a specific memory space. -class PinnedAllocation final : public Allocation { - public: - PinnedAllocation(HloPosition defining_position, MemorySpace memory_space, - std::optional chunk, - int64_t start_time, int64_t end_time); - - // Overridden methods - // - // Returns the original defining position. - HloPosition defining_position() const override; - int64_t earliest_available_time() const override { return start_time(); } - bool is_pinned_allocation() const override { return true; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return false; } - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - bool operator==(const PinnedAllocation& other) const; -}; - -// This class represents an allocation that is used to reserve a chunk of -// memory. If an HloPosition or an HloUse is colored in alternate memory, to -// make sure we are able to satisfy the coloring requirements, we reserve a -// chunk in the alternate memory before we start processing the buffers in -// sorted order. The reserved chunk serves as a fallback in case we are not able -// to satisfy the coloring requirements using the buffers in sorted order. -class ReservedAllocation final : public Allocation { - public: - ReservedAllocation(HloPosition defining_position, HeapSimulator::Chunk chunk, - int64_t start_time, int64_t end_time); - - // Overridden methods - // - // Returns the original defining position. - HloPosition defining_position() const override; - int64_t earliest_available_time() const override { return start_time(); } - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return true; } - bool is_mirrored_allocation() const override { return false; } - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - bool operator==(const ReservedAllocation& other) const; - - bool is_chunk_reserved_in_interval_tree() const { return reserved_; } - void mark_chunk_freed_in_interval_tree() { reserved_ = false; } - void mark_chunk_reserved_in_interval_tree() { reserved_ = true; } - - private: - // Indicates whether the chunk is still reserved in the interval_tree_. - bool reserved_; -}; - -// This class represents an allocation as a result of a single asynchronous -// data movement operation. The data movement operation is a copy, except when -// certain arguments are set, as described below. -// * CopyStart instructions are inserted after `copy_start_schedule_after`, -// while CopyDone instructions are inserted before -// `copy_done_schedule_before_time`. -// * When `sync_mem_op` is set, it points to a sync data movement instruction -// that we intend to turn into an asynchronous data movement operation. -// * When `async_mem_op_start` and `async_mem_op_done` are set, it indicates -// that an asynchronous data movement operation already exists, but MSA -// needs to appropriately schedule the operation. -// * If `sync_mem_op` is non-null, `async_mem_op_start` and `async_mem_op_done` -// must be null. -// * If are `async_mem_op_start` and `async_mem_op_done` are non-null, -// `sync_mem_op` must be null. -// * `source_operand_index` is the index of the operand in `sync_mem_op` -// (and consequently in the created `async_mem_op_start`) that represents the -// buffer being copied. This operand will be replaced with the actual -// producing instruction of the buffer (say, after tuples and GTEs are added). -class CopyAllocation final : public Allocation { - public: - CopyAllocation( - Allocation& prev_allocation, MemorySpace memory_space, - std::optional chunk, - int64_t copy_start_schedule_after_time, - int64_t copy_done_schedule_before_time, int64_t end_time, - std::optional cross_program_prefetch_index = std::nullopt, - HloInstruction* sync_mem_op = nullptr, - HloInstruction* async_mem_op_start = nullptr, - HloInstruction* async_mem_op_done = nullptr, - int64_t source_operand_index = 0); - - // Overridden methods - // - HloPosition defining_position() const override; - // Returns the time the buffer is first available to be used. For - // CopyAllocation, this is when the copy ends, which is - // copy_done_schedule_before. - int64_t earliest_available_time() const override; - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return true; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return false; } - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - const HloInstruction* sync_mem_op() const { return sync_mem_op_; } - bool operator==(const CopyAllocation& other) const; - - const Allocation& prev_allocation() { return prev_allocation_; } - Allocation& mutable_prev_allocation() { return prev_allocation_; } - - HloInstruction* copy_start() const { return copy_start_; } - HloInstruction* copy_done() const { return copy_done_; } - - void set_copy_start_schedule_after(int64_t copy_start_schedule_after); - void set_copy_done_schedule_before(int64_t copy_done_schedule_before); - int64_t copy_start_schedule_after() const; - int64_t copy_done_schedule_before() const; - - private: - Allocation& prev_allocation_; - // These variables define the scheduling boundaries where CopyStart and - // CopyDone can be scheduled. The earliest CopyStart can be scheduled is - // after copy_start_schedule_after_ and the latest CopyDone can be scheduled - // is before copy_done_schedule_before_. - int64_t copy_start_schedule_after_; - int64_t copy_done_schedule_before_; - HloInstruction* copy_start_ = nullptr; - HloInstruction* copy_done_ = nullptr; - // The sync data movement instruction that this copy is associated with. - HloInstruction* sync_mem_op_ = nullptr; - // The index of the operand in the async start instruction that should be - // replaced with the producing instruction. - int64_t source_operand_index_ = 0; -}; - -// This class represents an allocation resulting from a collection of -// asynchronous sliced copies that work together to copy a single tensor. -// -// Let the sliced allocation be represented as follows, and imagine that t3 -// is the time when the entire buffer [p0, p3) is available for use -// -// space -// ^ -// p3 | +-----------+ -// | | | -// p2 | +---+ | -// | | | -// p1 | +-------+ | -// | | | -// p0 | +-------+ -// +---|---|---|---|---|----> time -// t0 t1 t2 t3 t4 -// -// The PinnedAllocation underlying the SlicedCopyAllocation will use the -// following dimensions: -// - chunk = [p0, p3) -// - start time = t2 -// - earliest_available_time = t3 -// - end_time = t4 -class SlicedCopyAllocation final : public Allocation { - public: - // Full details about a slice in the sliced allocation. - struct SliceDetail { - std::string ToString() const; - std::tuple - ToTuple() const; - bool operator==(const SliceDetail& other) const; - - // Create the instructions to copy the slice. This method updates - // copy_start and copy_done. - absl::Status CreateAsyncSlice(const Shape& original_shape, - HloInstruction& producer, - HloComputation& parent); - - SliceDecision slice_decision; - int64_t copy_start_after_time = -1; - int64_t copy_done_before_time = -1; - HloInstruction* copy_start = nullptr; - HloInstruction* copy_done = nullptr; - }; - - // REQUIRES: - // - slice_decisions_sorted_by_exclusive_start_time.size() >= 2, otherwise, - // CopyAllocation should be used. - SlicedCopyAllocation( - const Allocation& prev_allocation, MemorySpace memory_space, - std::vector slice_decisions_sorted_by_exclusive_start_time, - int64_t copy_done_schedule_before_time, int64_t end_time, - const SlicedPrefetchOptions& sliced_prefetch_options, - absl::FunctionRef get_equivalent_s8_shape_fn, - HloInstruction* sync_mem_op = nullptr); - - // Overridden methods - // - HloPosition defining_position() const override; - // Returns the time the buffer is first available to be used. For - // SlicedCopyAllocation, this is when all copies have ended. - int64_t earliest_available_time() const override; - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return true; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return false; } - // MemorySpaceAssignment::Process() calls Process(const BitcastSplitFn& - // bitcast_split_fn) to create asynchronous slice copies, and a bitcast-concat - // call to glue the slices back together. - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - // Marks the allocation as needed. - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - const HloInstruction* sync_mem_op() const { return sync_mem_op_; } - bool operator==(const SlicedCopyAllocation& other) const; - - std::vector SliceOffsetsSortedByStartTime() const; - void AddDiffToAllSliceOffsets(int64_t diff); - // Used to update offsets and start times after repacking. - void ImportRepackedSliceData(const SlicedAllocationData& data); - const std::vector& slice_details_sorted_by_start_time() const; - std::vector& mutable_slice_details_sorted_by_start_time(); - HloInstruction* concat() const { return concat_; } - - private: - SlicedCopyAllocation() = delete; - - // Create an instruction to concatenate the slices. Populates concat_. - absl::Status CreateBitcastConcat(const Shape& shape, - absl::Span slices); - - Shape original_shape_to_slice_; - const Allocation& prev_allocation_; - // REQUIRES: - // - sorted_segments_[i].copy_start_after_time <= - // sorted_segments_[i+j].copy.start_after_time - // - sorted_segments_[i].copy_done_before_time <= - // sorted_segments_[i+j].copy.start_before_time - std::vector slice_details_sorted_by_exclusive_start_time_; - HloInstruction* concat_ = nullptr; - const SlicedPrefetchOptions& sliced_prefetch_options_; - absl::FunctionRef get_equivalent_s8_shape_fn_; - // The sync data movement instruction that this copy is associated with. - HloInstruction* sync_mem_op_ = nullptr; -}; - -// This class represents an allocation resulting from asynchronously prefetching -// a window buffer. When a tensor is placed in the default memory, we can -// prefetch the window buffer of the tensor to the alternate memory space. This -// is called window prefetching. -class WindowPrefetchedAllocation final : public Allocation { - public: - struct Options { - int64_t bytes = 0; - int64_t uid = 0; - int64_t alternate_memory_space = 0; - std::function - notify_operand_appended_fn = - [](const HloInstruction*, int64_t, int64_t) {}; - }; - - WindowPrefetchedAllocation(Allocation& prev_allocation, HloUse use, - const HeapSimulator::Chunk& chunk, - int64_t prefetch_start_schedule_after_time, - int64_t prefetch_done_schedule_before_time, - const Options& options); - - // Overridden methods - // - HloPosition defining_position() const override; - int64_t earliest_available_time() const override; - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return true; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return false; } - // MemorySpaceAssignment::Process() calls Process(const BitcastSplitFn& - // bitcast_split_fn) to create asynchronous window prefetches. - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - // Marks the allocation as needed. - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const WindowPrefetchedAllocation& other) const; - bool operator==(const Allocation& other) const override; - int64_t bytes() const { return bytes_; } - int64_t prefetch_start_schedule_after() const { - return prefetch_start_schedule_after_; - } - int64_t prefetch_done_schedule_before() const { - return prefetch_done_schedule_before_; - } - HloInstruction* prefetch() const { return prefetch_instruction_; } - - private: - // This method is called by Process(const BitcastSplitFn& bitcast_split_fn) to - // create window prefetch instructions. These instructions include a pair of - // async WindowPrefetch which is passed to the fusion. - absl::Status InsertWindowPrefetchInstruction( - HloInstruction* producing_instruction, HloInstruction* use_instruction, - HloComputation* computation); - - Options options_; - HloInstruction* prefetch_instruction_ = nullptr; - HloPosition defining_position_; - HloUse use_; - int64_t prefetch_start_schedule_after_; - int64_t prefetch_done_schedule_before_; - int64_t bytes_; -}; - -// An allocation that mirrors/duplicates another Allocation but does not reserve -// any new memory. We use a mirrored allocation for an HloPosition, when another -// allocation for a different HloPosition of the same HloValue exists, and we -// can read the value from that allocation instead of creating a new allocation. -// Mirrored allocations are used in the following cases: -// * The input to a while loop is in default memory, we prefetch the input in -// the while loop, for one use, but other uses in the while loop, after the -// prefetched uses, are served from default memory. In such a case, we serve -// those subsequent uses from a mirrored allocation of the input (in default -// memory) leading into the while loop, instead of performing an unnecessary -// evcition (in the while loop) of the prefetched instance. -// * If an argument to a conditional has been pinned to vmem for the duration of -// the conditional, we associate a mirrored allocation with the corresponding -// parameters in the conditional's branch computations. -class MirroredAllocation final : public Allocation { - public: - MirroredAllocation(const Allocation& original_allocation, int64_t time); - MirroredAllocation(HloPosition defining_position, - const Allocation& original_allocation, int64_t start_time, - int64_t end_time); - - // Overridden methods - // - // Returns the original defining position. - HloPosition defining_position() const override; - int64_t earliest_available_time() const override { return start_time(); } - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return true; } - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - bool operator==(const MirroredAllocation& other) const; - const Allocation& original_allocation() const { return original_allocation_; } - - HeapSimulator::Chunk chunk() const final { - return original_allocation_.chunk(); - } - std::optional maybe_chunk() const final { - return original_allocation_.maybe_chunk(); - } - - private: - const std::optional defining_position_; - const Allocation& original_allocation_; -}; - -// An allocation in default memory space that is defined in the parent -// computation. If a value has a copy in the default memory space in the -// parent computation, we don't need to evict this buffer in a while loop. -class ParentAllocation final : public Allocation { - public: - ParentAllocation(const Allocation& original_allocation, - HloInstruction* calling_instruction, HloPosition position, - int64_t time); - - // Overridden methods - // - // Returns the original defining position. - HloPosition defining_position() const override; - int64_t earliest_available_time() const override { return start_time(); } - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return false; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return false; } - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override; - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - bool operator==(const ParentAllocation& other) const; - - private: - const Allocation& original_allocation_; - HloInstruction* calling_instruction_; -}; - -// An allocation representing scoped alternate memory. -class ScopedAllocation final : public Allocation { - public: - // is_post_module is true if the allocation is for a scoped allocation that - // is used after the module. - ScopedAllocation(HeapSimulator::Chunk chunk, int64_t allocation_time, - HloInstruction* defining_instruction, bool is_post_module); - - // Overridden methods - HloPosition defining_position() const override; - int64_t earliest_available_time() const override { return start_time(); } - bool is_pinned_allocation() const override { return false; } - bool is_copy_allocation() const override { return false; } - bool is_sliced_copy_allocation() const override { return false; } - bool is_window_prefetched_allocation() const override { return false; } - bool is_scoped_allocation() const override { return true; } - bool is_reserved_allocation() const override { return false; } - bool is_mirrored_allocation() const override { return false; } - absl::Status Process(const BitcastSplitFn& bitcast_split_fn, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) override; - absl::Status PostProcess() override { return absl::OkStatus(); } - void MarkIfNeeded(absl::flat_hash_set& needed_allocations) - const override; - void MarkNeeded(absl::flat_hash_set& needed_allocations) - const override; - std::string ToString() const override; - bool operator==(const Allocation& other) const override; - - // New non-virtual methods - bool operator==(const ScopedAllocation& other) const; - bool is_post_module() const { return is_post_module_; } - - private: - bool is_post_module_; -}; - -// A class with some utility functions that are useful in debugging. -struct AllocationSequenceDebugging { - // Developers can call this method to log all the allocations in alternate - // memory, at a given instruction time. - // - // REQUIRED: - // - This method is intended to be called before MSA modifies the HloModule. - static void LogAltMemAllocationsAt(const AllocationSequence& allocations, - int64_t time); -}; - -} // namespace xla::memory_space_assignment - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALLOCATION_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc deleted file mode 100644 index 418658d2a20cce..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc +++ /dev/null @@ -1,300 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/allocation.h" - -#include -#include -#include - -#include -#include -#include "absl/strings/string_view.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.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/testlib/hlo_hardware_independent_test_base.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/shape_util.h" -#include "xla/xla_data.pb.h" - -namespace xla::memory_space_assignment { -namespace { - -class AllocationTest : public HloHardwareIndependentTestBase { - protected: - void RunAnalysis(HloModule* module, - const std::vector& inst_names, - std::unique_ptr& live_range, - std::unique_ptr& alias_analysis) { - HloSchedule schedule(module); - HloInstructionSequence sequence; - for (auto name : inst_names) { - sequence.push_back(FindInstruction(module, name)); - } - schedule.set_sequence(module->entry_computation(), sequence); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(alias_analysis, - HloAliasAnalysis::Run(module, &alias_info)); - ASSERT_OK_AND_ASSIGN(live_range, - HloLiveRange::Run(schedule, *alias_analysis, - module->entry_computation())); - } -}; - -TEST_F(AllocationTest, CopyAllocationProcessSimple) { - absl::string_view hlo_string = R"( -HloModule module - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p1_negate = f32[2,3]{1,0} negate(p1) - add = f32[2,3]{1,0} add(p0, p1_negate) - ROOT tuple = tuple(add, p0) -} - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - std::unique_ptr hlo_live_range; - std::unique_ptr alias_analysis; - RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "add", "tuple"}, - hlo_live_range, alias_analysis); - // HloComputation* computation = module->entry_computation(); - HloInstruction* add = FindInstruction(module.get(), "add"); - HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); - - HeapSimulator::Chunk p1_negate_chunk = - HeapSimulator::Chunk::FromOffsetSize(0, 24); - - PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, - MemorySpace::kDefault, p1_negate_chunk, - /*start_time=*/0, - /*end_time=*/5); - CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kAlternate, - std::nullopt, - /*copy_start_schedule_after_time=*/2, - /*copy_done_schedule_before_time=*/3, - /*end_time=*/5, std::nullopt, - /*sync_mem_op=*/nullptr); - - // Use the correct instruction and operand numbers for the add instruction - copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add - BitcastSplitFn split_fn = nullptr; - ASSERT_OK( - copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); - - // Check copy_start and copy_done instructions. - HloInstruction* copy_start = copy_allocation.copy_start(); - ASSERT_NE(copy_start, nullptr); - EXPECT_EQ(copy_start->opcode(), HloOpcode::kCopyStart); - EXPECT_EQ(copy_start->operand(0), p1_negate); - - HloInstruction* copy_done = copy_allocation.copy_done(); - ASSERT_NE(copy_done, nullptr); - EXPECT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); - EXPECT_EQ(copy_done->operand(0), copy_start); - - // Check that uses are updated. - EXPECT_EQ(add->operand(1), copy_done); - - // Check defining position - EXPECT_EQ(copy_allocation.defining_position().instruction, copy_done); -} - -TEST_F(AllocationTest, EvictedSplitShape) { - absl::string_view hlo_string = R"( -HloModule module - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p1_negate = f32[2,3]{1,0:S(1)SC(0:1)} negate(p1) - add = f32[2,3]{1,0} add(p0, p1_negate) - ROOT tuple = tuple(add, p0) -} - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - std::unique_ptr hlo_live_range; - std::unique_ptr alias_analysis; - RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "add", "tuple"}, - hlo_live_range, alias_analysis); - // HloComputation* computation = module->entry_computation(); - HloInstruction* add = FindInstruction(module.get(), "add"); - HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); - - HeapSimulator::Chunk p1_negate_chunk = - HeapSimulator::Chunk::FromOffsetSize(0, 24); - - PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, - MemorySpace::kAlternate, p1_negate_chunk, - /*start_time=*/0, - /*end_time=*/5); - CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kDefault, - std::nullopt, - /*copy_start_schedule_after_time=*/2, - /*copy_done_schedule_before_time=*/3, - /*end_time=*/5, std::nullopt, - /*sync_mem_op=*/nullptr); - - // Use the correct instruction and operand numbers for the add instruction - copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add - BitcastSplitFn split_fn = nullptr; - ASSERT_OK( - copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); - - // Check copy_start and copy_done instructions. - HloInstruction* copy_start = copy_allocation.copy_start(); - ASSERT_NE(copy_start, nullptr); - EXPECT_EQ(copy_start->opcode(), HloOpcode::kCopyStart); - EXPECT_EQ(copy_start->operand(0), p1_negate); - - HloInstruction* copy_done = copy_allocation.copy_done(); - ASSERT_NE(copy_done, nullptr); - EXPECT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); - EXPECT_EQ(copy_done->operand(0), copy_start); - EXPECT_EQ(copy_done->shape().layout().split_configs().size(), 0); - - // Check that uses are updated. - EXPECT_EQ(add->operand(1), copy_done); - - // Check defining position - EXPECT_EQ(copy_allocation.defining_position().instruction, copy_done); -} - -TEST_F(AllocationTest, CopyAllocationProcessReplaceSyncSlice) { - absl::string_view hlo_string = R"( -HloModule module - -ENTRY entry { - p0 = f32[1,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p1_negate = f32[2,3]{1,0} negate(p1) - slice = f32[1,3]{1,0} slice(p1_negate), slice={[0:1], [0:3]} - add = f32[1,3]{1,0} add(p0, slice) - ROOT tuple = tuple(add, p0) -} - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - std::unique_ptr hlo_live_range; - std::unique_ptr alias_analysis; - RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "slice", "add", "tuple"}, - hlo_live_range, alias_analysis); - // HloComputation* computation = module->entry_computation(); - HloInstruction* add = FindInstruction(module.get(), "add"); - HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); - HloInstruction* slice = FindInstruction(module.get(), "slice"); - - HeapSimulator::Chunk p1_negate_chunk = - HeapSimulator::Chunk::FromOffsetSize(0, 24); - - PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, - MemorySpace::kAlternate, p1_negate_chunk, - /*start_time=*/0, - /*end_time=*/5); - CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kAlternate, - std::nullopt, - /*copy_start_schedule_after_time=*/2, - /*copy_done_schedule_before_time=*/3, - /*end_time=*/5, std::nullopt, - /*sync_mem_op=*/slice); - - // Use the correct instruction and operand numbers for the add instruction - copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add - BitcastSplitFn split_fn = nullptr; - ASSERT_OK( - copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); - - // Check copy_start and copy_done instructions. - HloInstruction* slice_start = copy_allocation.copy_start(); - ASSERT_NE(slice_start, nullptr); - EXPECT_EQ(slice_start->opcode(), HloOpcode::kAsyncStart); - EXPECT_EQ(slice_start->operand(0), p1_negate); - - HloInstruction* slice_done = copy_allocation.copy_done(); - ASSERT_NE(slice_done, nullptr); - EXPECT_EQ(slice_done->opcode(), HloOpcode::kAsyncDone); - EXPECT_EQ(slice_done->operand(0), slice_start); - - // Check the shapes. - EXPECT_EQ(slice_done->shape(), slice->shape()); - - // Check that uses are updated. - EXPECT_EQ(add->operand(1), slice_done); - - // Check defining position - EXPECT_EQ(copy_allocation.defining_position().instruction, slice_done); -} - -TEST_F(AllocationTest, SkipTupleReconstructionForAsyncCollective) { - absl::string_view hlo_string = R"( -HloModule module - -ENTRY entry { - new_buffer = f32[2,3]{1,0} parameter(0) - cp-start = (f32[2,3]{1,0}, f32[2,3]{1,0}, u32[], u32[]) collective-permute-start(new_buffer), channel_id=1, source_target_pairs={{0,1}} - cp-done = f32[2,3]{1,0} collective-permute-done(cp-start) - ROOT tuple = tuple(cp-done) -} - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - std::unique_ptr hlo_live_range; - std::unique_ptr alias_analysis; - RunAnalysis(module.get(), {"new_buffer", "cp-start", "cp-done", "tuple"}, - hlo_live_range, alias_analysis); - - HloInstruction* cp_start = FindInstruction(module.get(), "cp-start"); - HloInstruction* cp_done = FindInstruction(module.get(), "cp-done"); - HloInstruction* new_buffer = FindInstruction(module.get(), "new_buffer"); - - HeapSimulator::Chunk chunk = HeapSimulator::Chunk::FromOffsetSize(0, 24); - - PinnedAllocation pinned(HloPosition{new_buffer, {}}, MemorySpace::kAlternate, - chunk, 0, 5); - pinned.AddUse(HloUse{cp_done, 0, {0}}); - - BitcastSplitFn split_fn = nullptr; - ASSERT_OK(pinned.Process(split_fn, *hlo_live_range, *alias_analysis)); - - EXPECT_EQ(cp_done->operand(0), cp_start); -} - -TEST_F(AllocationTest, MirroredAllocationDelegatesChunkToOriginalAllocation) { - HloComputation::Builder builder("entry"); - HloInstruction* p0 = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeShape(F32, {2, 3}), "p0")); - PinnedAllocation original_allocation( - HloPosition{p0, {}}, MemorySpace::kAlternate, - HeapSimulator::Chunk::FromOffsetSize(-1, 64), - /*start_time=*/0, /*end_time=*/5); - MirroredAllocation mirrored_allocation(original_allocation, /*time=*/2); - original_allocation.set_offset(128); - - const Allocation& alloc = mirrored_allocation; - ASSERT_TRUE(alloc.maybe_chunk().has_value()); - EXPECT_EQ(alloc.maybe_chunk()->offset, 128); - EXPECT_EQ(alloc.chunk().offset, 128); -} - -} // namespace -} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc deleted file mode 100644 index 010c5737311e18..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc +++ /dev/null @@ -1,112 +0,0 @@ -/* Copyright 2025 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/allocation_value.h" - -#include - -#include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" - -namespace xla::memory_space_assignment { - -std::string AllocationValue::ToString() const { - std::string out = absl::StrCat("computation = ", computation()->name()); - absl::StrAppend( - &out, (requires_contiguous_allocation_ ? " (contiguous alloc)" : "")); - absl::StrAppend(&out, "\n position:\n"); - absl::StrAppend(&out, " ", position_.ToString(), "\n"); - absl::StrAppend(&out, " uses:\n"); - for (const Use& use : uses_) { - absl::StrAppend(&out, " ", use.hlo_use.ToString(), "\n"); - } - return out; -} - -std::string AllocationValue::ToShortString() const { - return absl::StrCat("computation = ", computation()->name(), - ", position = ", position_.ToString(), - ", value = ", value_->ToShortString(), - (requires_contiguous_allocation_ ? " (cont alloc)" : "")); -} - -std::string AllocationRequest::ToString() const { - return absl::StrJoin( - {absl::StrCat("size: ", size), - absl::StrCat("inclusive_start_time: ", inclusive_start_time, - "; end_time: ", end_time, - "; latest_prefetch_time: ", latest_prefetch_time, - "; required_copy_allocation_latest_time: ", - required_copy_allocation_latest_time), - absl::StrCat("prefer_no_copy_alternate_mem_allocation: ", - prefer_no_copy_alternate_mem_allocation, - "; allow_no_copy_alternate_mem_allocation: ", - allow_no_copy_alternate_mem_allocation, - "; require_no_copy_alternate_mem_allocation: ", - require_no_copy_alternate_mem_allocation, - "; require_copy_allocation: ", require_copy_allocation), - absl::StrCat("allow_prefetch: ", allow_prefetch), - absl::StrCat("earliest_prefetch_time: ", - earliest_prefetch_time.has_value() - ? absl::StrCat(*earliest_prefetch_time) - : "nullopt"), - absl::StrCat("preferred_prefetch_time: ", - preferred_prefetch_time.has_value() - ? absl::StrCat(*preferred_prefetch_time) - : "nullopt"), - absl::StrCat( - "preferred_offset: ", - preferred_offset - ? absl::StrCat(preferred_offset->offset, "; allocations: ", - preferred_offset->allocations.size()) - : "nullptr"), - absl::StrCat("use: ", use ? use->hlo_use.ToString() : "nullptr"), - absl::StrCat("allocation_value: ", allocation_value - ? allocation_value->ToString() - : "nullptr"), - absl::StrCat("allocation_value_to_update: ", - allocation_value_to_update - ? allocation_value_to_update->ToString() - : "nullptr"), - absl::StrCat("all_use_times: ", all_use_times.size(), " elements"), - absl::StrCat("required_copy_allocation_for: ", - required_copy_allocation_for - ? required_copy_allocation_for->ToString() - : "nullptr"), - absl::StrCat("required_copy_for_slice: ", required_copy_for_slice), - absl::StrCat("only_extend_existing_allocation: ", - only_extend_existing_allocation), - absl::StrCat("processed_allocation_values: ", - processed_allocation_values.size(), " elements"), - absl::StrCat("no_copy_chunk_inclusive_start_time: ", - no_copy_chunk_inclusive_start_time.has_value() - ? absl::StrCat(*no_copy_chunk_inclusive_start_time) - : "nullopt"), - absl::StrCat( - "require_start_colored_in_alternate_memmory: ", - require_start_colored_in_alternate_memory, - "; require_end_colored_in_alternate_memory: ", - require_end_colored_in_alternate_memory, - "; require_start_colored_in_default_memory: ", - require_start_colored_in_default_memory, - "; require_end_colored_in_default_memory: ", - require_end_colored_in_default_memory, - "; strict_timing: ", strict_timing, - "; fail_on_unsatisfied_override: ", fail_on_unsatisfied_override, - "; is_prefetch_override: ", is_prefetch_override)}, - "\n"); -} - -} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h deleted file mode 100644 index f3e8f7d3dca438..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h +++ /dev/null @@ -1,342 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALLOCATION_VALUE_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALLOCATION_VALUE_H_ - -#include -#include -#include -#include - -#include "absl/container/flat_hash_set.h" -#include "absl/strings/str_cat.h" -#include "absl/types/span.h" -#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 { -// AllocationValue is used to break up HloValues for each non-trivial position -// (trivial positions are considered Tuple, GetTupleElement, and Bitcast). An -// HloValue may include positions and uses that alias with each other across -// multiple computations. We use this class to break these HloValues such that -// every AllocationValue has one position (that may alias with other -// AllocationValues). The uses field of the AllocationValue contains only the -// direct uses of the AllocationValue's position. -// -// For example, consider the following HLO snippet: -// -// Body { -// body_param = (f32[4,3]{1,0}, f32[]) parameter(0) -// get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element(body_param), -// index=0 -// ... -// ROOT tuple = (f32[4,3]{1,0}, f32[]) tuple(get-tuple-element.3, ...) -// } -// -// Cond { -// cond_param = (f32[4,3]{1,0}, f32[]) parameter(0) -// ... -// } -// -// add.4 = f32[4,3]{1,0} add(...) -// tuple.1 = (f32[4,3]{1,0}, f32[]) tuple(add.4, ...) -// while = (f32[4,3]{1,0}, f32[]) while(tuple.1), body=Body, condition=Cond -// get-tuple-element.5 = f32[4,3]{1,0} get-tuple-element(while), index=0 -// add.5 = f32[4,3]{1,0} add(get-tuple-element.5, ...) -// -// This contains an HloValue that looks like the following: -// positions: -// add.4 -// body_param {0} -// get-tuple-element.3 -// tuple {0} -// cond_param {0} -// tuple.1 {0} -// while {0} -// get-tuple-element.5 -// uses: -// tuple.1, operand 0 -// tuple, operand 0 -// while, operand 0 {0} -// add.5, operand 0 -// -// We break this HloValue up into the following AllocationValues for each -// non-trivial position: -// AllocationValue1: computation = Entry -// position: -// add.4 -// uses: -// while, operand 0 {0} -// AllocationValue2: computation = Cond -// position: -// cond_param {0} -// uses: -// AllocationValue3: computation = Body -// position: -// body_param {0} -// uses: -// add.1, operand 0 -// tuple, operand 0 -// AllocationValue4: computation = Entry -// position: -// while {0} -// uses: -// add.5, operand 0 -class AllocationValue { - public: - // This data structure wraps an HloUse and adds additional metadata that are - // useful for allocation. - struct Use { - // The wrapped HloUse object. - HloUse hlo_use; - // The logical time this use is scheduled. - int64_t time; - // All the positions where this use aliases with. The aliased positions - // must get the same allocation. - std::vector aliases; - // A synchronous memory operation that feeds this use. - // TODO(mehrdadk): extend this to support multiple sync data movement - // operands. - HloInstruction* sync_mem_op_operand = nullptr; - - bool operator==(const Use& other) const { - return hlo_use == other.hlo_use && time == other.time && - aliases == other.aliases; - } - - template - friend H AbslHashValue(H h, const Use& s) { - return H::combine(std::move(h), s.hlo_use, s.time, s.aliases); - } - }; - - AllocationValue(const HloValue* value, const HloPosition& position, - int64_t size) - : value_(value), - position_(position), - size_(size), - requires_contiguous_allocation_(false), - split_shape_(std::nullopt) {} - - const HloPosition& position() const { return position_; } - int64_t size() const { return size_; } - void set_size(int64_t size) { size_ = size; } - const std::vector& uses() const { return uses_; } - std::vector& uses() { return uses_; } - const HloValue* value() const { return value_; } - const HloComputation* computation() const { - return position().instruction->parent(); - } - AllocationSequence* mutable_allocation_sequence() { - return &allocation_sequence_; - } - const AllocationSequence* allocation_sequence() const { - return &allocation_sequence_; - } - - // Sets/gets whether this AllocationValue requires allocating it - // contiguously throughout its live range (without any copies). - bool requires_contiguous_allocation() const { - return requires_contiguous_allocation_; - } - void set_requires_contiguous_allocation(bool requires_contiguous_allocation) { - requires_contiguous_allocation_ = requires_contiguous_allocation; - } - - void AddUse(const HloUse& use, int64_t use_time) { - uses_.push_back({use, use_time, {}}); - } - - void set_split_shape(const Shape& split_shape) { split_shape_ = split_shape; } - std::optional mutable_split_shape() { return split_shape_; } - - std::string ToString() const; - std::string ToShortString() const; - - private: - const HloValue* value_; - HloPosition position_; - int64_t size_; - // If true, there must be a contiguous allocation for this buffer without - // any copies. - bool requires_contiguous_allocation_; - std::vector uses_; - AllocationSequence allocation_sequence_; - - // If present, indicates the newly split shape. - std::optional split_shape_; -}; - -// A data structure we use to associate Allocation objects that are aliased -// and must get the same offset. -struct AliasedOffset { - int64_t offset; - absl::flat_hash_set allocations; -}; - -// An allocation request for a use segment. A use segment is the time segment -// between the definition and the first use, and the time segment between the -// uses of a buffer. For example, the time between the definition and Use1, is -// the first segment, and the time between Use1 and Use2 is the second segment -// and so on: -// -// +------+----------+-------+ -// / \ \ \ -// / v v v -// Def Use1 Use2 Use3 -// <----------> <--------> <-----> -// Segment Segment Segment -// -// start_time and end_time are the start and end logical times of the segment. -// use_times is a sorted sequence of the times of all uses. -// latest_prefetch_time is the latest time we can schedule the CopyDone for a -// prefetch. -// If allow_no_copy_alternate_mem_allocation is false, an eviction is forced. -// If earliest_prefetch_time is set, prefetches cannot start before this -// value. -// -// In case we are trying to replace synchronous copies, and for example Use2 -// is a replaceable sync copy candidate, we now skip Use2 and segments will be -// between Def, Use1, Use2.1, Use2.2, Use3: -// +------+----------+-------------------+ -// / \ \ \ -// / v v v -// Def Use1 Use2(Sync Copy) Use3 -// | | \ \ | -// | | v v | -// | | Use2.1 Use2.2 | -// |<---------->|<---------->|<------->|<----->| -// | Segment | Segment | Segment |Segment| - -struct AllocationRequest { - int64_t inclusive_start_time; - int64_t end_time; - int64_t latest_prefetch_time; - // See the comment for require_copy_allocation - int64_t required_copy_allocation_latest_time; - int64_t size; - bool prefer_no_copy_alternate_mem_allocation; - bool allow_no_copy_alternate_mem_allocation; - bool require_no_copy_alternate_mem_allocation; - // If true, indicates we are requiring a copy allocation between def and - // use, that finishes by required_copy_allocation_latest_time. - // required_copy_allocation_for is a synchronous copy instruction that will - // be removed, if we are successful in adding the copy allocation. - bool require_copy_allocation; - bool allow_prefetch; - std::optional earliest_prefetch_time; - std::optional preferred_prefetch_time; - AliasedOffset* preferred_offset; - const AllocationValue::Use* use; - AllocationValue* allocation_value; - absl::Span all_use_times; - // See the comment for require_copy_allocation - HloInstruction* required_copy_allocation_for; - // If the required copy in require_copy_allocation is only for a slice of - // the allocation_value - bool required_copy_for_slice; - // The resulting Allocation will be added to the AllocationSequence of - // allocation_value_to_update. We only expect allocation_value_to_update to - // be different from allocation_value in the case of a synchronous memory - // operation conversion to asynchronous, otherwise, they should be the same. - AllocationValue* allocation_value_to_update; - // No new Allocation is needed to be created and we will only extend an - // existing one. - bool only_extend_existing_allocation; - // Data structure that contains the options for making window prefetched - // allocations. - const WindowPrefetchedAllocation::Options* window_prefetch_options = nullptr; - // Previously processed AllocationValues, with the same parent HloValue as the - // request. - absl::Span processed_allocation_values; - // An optional override starting time for the placement of a chunk on the MSA - // heap, for a no-copy allocation (see - // MsaAlgorithm::AllocateInAlternateMemoryNoCopy() for more details). - // - // Note, this override is used when an aliased AllocationValue has already - // done some of the heap allocation for us. So this request picks up where it - // left off. - std::optional no_copy_chunk_inclusive_start_time; - // An optional override for the shape for PrefetchContext::full_shape, which - // is used to calculate space needed for the prefetch. - std::optional shape_override; - // Indicates if the AllocationRequest start time (definition time) has an - // alternate memory color requirement. - bool require_start_colored_in_alternate_memory = false; - // Indicates if the AllocationRequest end time (use time) has an alternate - // memory color requirement. - bool require_end_colored_in_alternate_memory = false; - // Indicates if the AllocationRequest start time (definition time) has a - // default memory color requirement. - bool require_start_colored_in_default_memory = false; - // Indicates if the AllocationRequest end time (use time) has a default - // memory color requirement. - bool require_end_colored_in_default_memory = false; - // If true, forces the prefetch to start at the exact preferred_prefetch_time - // without searching outwards. - bool strict_timing = false; - // If true, any failure to satisfy this override results in a fatal error. - bool fail_on_unsatisfied_override = false; - // Indicates if this request has a prefetch override. - bool is_prefetch_override = false; - - std::string ToString() const; -}; - -// Result of an allocation, prefetch, eviction etc. request. The result is -// either kSuccess or a bitwise OR of one or more failures. The values are -// unique powers of two. To check if a result contains a particular failure, -// use the result_is method. To add a new failure to a result, use the -// result_mark method. -enum class AllocationResult { - // Successful allocation. - kSuccess = 0, - // Allocation failed because we ran out of alternate memory. - kFailOutOfMemory = 1, - // A no-copy allocation couldn't be performed because the previous - // allocation wasn't in the alternate memory space. - kFailPrevAllocationNotInAlternateMem = 2, - // A no-copy allocation couldn't be performed because the live range was too - // long. - kFailLiveRangeTooLong = 4, - // A prefetching couldn't be performed because the live range was too short. - kFailLiveRangeTooShort = 8, - // Ran out of outstanding asynchronous copy limit either during prefetching - // or eviction. - kFailOutOfAsyncCopies = 16, - // A prefetching couldn't be performed because the asynchronous copy - // resource was violated. - kFailViolatesAsyncCopyResource = 32, - // An allocation failure happened that requires uncommitting all the pending - // allocations. Usually this is due to a situation requiring an eviction but - // the eviction couldn't be performed. - kFailRequiresUncommit = 64, - // For prefetching, indicates that all slices have the same start time, in - // which case, we fallback to an unsliced solution. - kAllSlicesHaveTheSameStartTime = 128, - // There were conflicting preferred offsets. - kFailConflictingPreferredOffsets = 256, - // Could not replace the synchronous data movement instruction (e.g., kCopy, - // kSlice) with an asynchronous one - kFailSyncDataMoveReplacement = 512 -}; - -} // namespace memory_space_assignment -} // namespace xla -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_ALLOCATION_VALUE_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.cc b/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.cc deleted file mode 100644 index 0e241733f12044..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.cc +++ /dev/null @@ -1,650 +0,0 @@ -/* Copyright 2020 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -/* -BestFitRepacker Algorithm - -Preliminary terminology/ideas - * Consider a sliced AllocationBlock B with S slices. - - Let s_{B,i} be the size of slice at the ith smallest offset. - - Let s_B be the sum(s_{B,i}, for all i). - - t_{B,i} be the ith earliest start time for a slice in B. (Note, an - AllocationBlock has not determined which slice offset maps to which slice - start time.) - - Let e_B be the allocation end time for all slices. - - We defined the following: - - B's full buffer interval - * Defined by the buffer interval spanning time [s_{B,0}, e_B], with size - sum(s_{B, i}, for all i). - * In the illustration below, B's full buffer interval occupies W+X+Y+Z - * In more colloquial language, the smallest rectangle that can bound - a sliced allocation. - - B's minimum buffer interval - * Defined by the buffer interval spanning the time [t_{B,S-1}, e_B], - with size s_B. - * In the illustration below, B's minimum buffer interval occupies W. - * In more colloquial language, the sliced allocation once all slices are - allocated. - - Potential placement of B - * A potential placement of B occupies W+X. Note other placements are - possible, e.g., starting the slice at offset o at time t_{B,0} or - t_{B,2}. - - Illustration of B with S=3: - space - ^ - o+s_B | +-------------------------+------------+ - | | Y | | - o+s_{B,0}+s_{B,1} | +-------------------------+ | - | | | W | - o+s_{B,0} | +------------+ X | | - | | Z | | | - o | +------------+------------+------------+ - | - +--|------------|------------|------------|-> time - t_{B,0} t_{B,1} t_{B,2} e_B - - * For a non-sliced AllocationBlock the full buffer interval == the miniumum - buffer interval == potential placement (given starting offset o). - -Step 1: Data structure construction (ImportAllocationBlocks()) - * allocation_blocks_: Contains the input AllocationBlocks to repack. If we - find a valid repacking, we update these AllocationBlocks with repacking - placement data. - * buffer_intervals_: Data structure owned by GlobalDecreasingSizeBestFitHeap. - We use it to map each each AllocationBlock to its minimum buffer - interval. - * full_buffer_interval_map_: Maps each AllocationBlock to its minimum buffer - interval. - * sliced_buffer_interval_map_: Maps each AllocationBlock to a - SlicedBufferInterval that wraps the corresponding full buffer interval - stored in full_buffer_interval_map_. - -Step 2: Sort buffers (GetSortedBufferIntervals()) - * This step prioritizes the order in which we will try place each buffer in - the repacking. - * GetSortedBufferIntervals() sorts the buffer intervals in buffer_intervals_. - The DefaultBufferIntervalCompare() maps each buffer in buffer_intervals to - its full buffer interval, and uses the full buffer interval's - properties/dimensions for sorting. - -Step 3: Find and commit buffer chunks (FindAndCommitChunks()) - * We iterate through AllocationBlocks (in the sorted order from Step 2), - finding a location for them. - * When we try to find a placement for an AllocationBlock B, we also find - locations for its colocations. Colocations are done in tandem with B because - B cannot be committed to a location if the same offset does not work for its - colocations. - * Colocations are handled as follows: - - If a colocation is not sliced, MakeFreeChunks() will make sure we do not - consider any placement for B that does not also accomodate the colocation. - - If a colocation is sliced, MakeFreeChunks() will be inadequate for - enforcing a placement for the colocation. The problem is that during - placement, we have not yet determined the mapping between offsets and - slice times for a sliced colocation. So, we have 2 options. - Option 1) Overestimate the locations of free space available for B, - knowing that some may not work when we try to place sliced - colocations. - Option 2) Underestimate the locations of free space available for B, - knowing this will cause us not to choose some valid placement - options for B. - We have chosen option 1. To compensate for the free space over estimation, - every time we find a place for B, we must explicitly check that each of - its sliced colocations will also fit. - -Step 4: Check if the repacking fits in the heap size - * If the repacking does not fit in the heap size, the repacking fails. - -Step 5: Update AllocationBlocks with the repacking placements - * We update the offset and repacked_slicing fields of an AllocationBlock. - Callers extract that data to get the repacking locations. -*/ - -#include "xla/service/memory_space_assignment/best_fit_repacker.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/functional/any_invocable.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/log/vlog_is_on.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" -#include "absl/types/span.h" -#include "xla/comparison_util.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" - -namespace xla { -namespace { - -bool IsSliced(const AllocationBlock* block) { - return block->original_slice_data.has_value(); -} - -template -std::vector SortAllocationBlocks(const T& container) { - std::vector result; - result.insert(result.end(), container.begin(), container.end()); - absl::c_sort( - result, [](const AllocationBlock* lhs, const AllocationBlock* rhs) { - return std::make_tuple(lhs->inclusive_start_time, lhs->end_time, - lhs->initial_offset, lhs->size) < - std::make_tuple(rhs->inclusive_start_time, rhs->end_time, - rhs->initial_offset, rhs->size); - }); - - return result; -} - -const SlicedAllocationData* GetSlicedAllocationDataPointer( - const std::optional& sliced_allocation_data) { - if (!sliced_allocation_data.has_value()) { - return nullptr; - } - return &(*sliced_allocation_data); -} - -// A slice-aware best-fit repacker. -class BestFitRepacker - : public GlobalDecreasingSizeBestFitHeap { - public: - BestFitRepacker( - const memory_space_assignment::MemorySpaceAssignmentBestFitRepacker:: - BestFitRepackOptions& options, - SliceTimePermutationIterator::Ty slice_time_permutation_iterator_type, - int64_t max_size, int64_t alignment) - : GlobalDecreasingSizeBestFitHeap( - alignment, kCustom, - (options.buffer_interval_compare ? options.buffer_interval_compare - : DefaultBufferIntervalCompare()), - slice_time_permutation_iterator_type), - validate_(options.validate), - max_size_(max_size) {} - - // Initialize our basic data structures: allocation_blocks_, - // buffer_intervals_, full_buffer_interval_map_, and - // sliced_buffer_interval_map_. - void ImportAllocationBlocks(absl::Span allocations) { - allocation_blocks_ = allocations; - // We loop through allocation_blocks_ once to build - // full_buffer_interval_map_, with colocations fully specified. - for (AllocationBlock* allocation_block : allocation_blocks_) { - // Check if any of the colocations are already added to buffer_intervals_. - CHECK_NE(allocation_block->next_colocated, nullptr); - if (full_buffer_interval_map_.contains(allocation_block)) { - VLOG(4) << "Skipping " << allocation_block->ToString() - << " because it was already added to full_buffer_interval_map_ " - "while processing its colocations."; - continue; - } - full_buffer_interval_map_.insert( - std::make_pair(allocation_block, - BufferInterval{allocation_block, - allocation_block->size, - allocation_block->inclusive_start_time, - allocation_block->end_time, - {}, - /*need_allocation=*/true})); - for (AllocationBlock* colocated = allocation_block->next_colocated; - colocated != allocation_block; - colocated = colocated->next_colocated) { - full_buffer_interval_map_.insert(std::make_pair( - colocated, BufferInterval{colocated, - colocated->size, - colocated->inclusive_start_time, - colocated->end_time, - {}, - /*need_allocation=*/false})); - full_buffer_interval_map_.at(allocation_block) - .colocations.push_back(colocated); - } - } - - // Now that full_buffer_interval_map_ has full colocation specifications, - // we loop through allocation_blocks_ again to build - // sliced_buffer_interval_map_ and buffer_intervals_. Also note, at this - // point we will not longer add or remove items from - // full_buffer_interval_map_. This is important because - // sliced_buffer_interval_map_ will retain pointers to BufferIntervals in - // full_buffer_interval_map_. - for (AllocationBlock* allocation_block : allocation_blocks_) { - BufferInterval& full_buffer_interval = - full_buffer_interval_map_[allocation_block]; - SlicedBufferInterval& sliced_buffer_interval = - sliced_buffer_interval_map_ - .insert(std::make_pair( - allocation_block, SlicedBufferInterval::CreateMutableInterval( - full_buffer_interval))) - .first->second; - if (IsSliced(allocation_block)) { - const SlicedAllocationData& original_slice_data = - allocation_block->original_slice_data.value(); - CHECK(!original_slice_data.slices_sorted_by_offset.empty()); - - sliced_buffer_interval.Slice(original_slice_data.SizesSortedByOffset()); - sliced_buffer_interval.UpdateInclusiveSliceStartTimes( - original_slice_data.SortedInclusiveStartTimes()); - } - - // We use buffer_intervals_ to store the minimum buffer interval for - // allocation_block. See the algorithm description (at the head of this - // file) for more details. - buffer_intervals_[allocation_block] = - sliced_buffer_interval.IntervalForMakeFreeChunks( - sliced_buffer_interval.num_slices() - 1); - } - - CHECK_EQ(allocation_blocks_.size(), buffer_intervals_.size()); - CHECK_EQ(allocation_blocks_.size(), full_buffer_interval_map_.size()); - CHECK_EQ(allocation_blocks_.size(), sliced_buffer_interval_map_.size()); - - VLOG(2) << [&]() -> std::string { - int sliced_blocks = 0; - int colocation_sets = 0; - int colocation_sets_with_multiple_sliced_blocks = 0; - absl::flat_hash_set seen_blocks; - - for (const auto& allocation_and_buffer_interval : buffer_intervals_) { - const AllocationBlock* block = allocation_and_buffer_interval.first; - const BufferInterval& min_buffer_interval = - allocation_and_buffer_interval.second; - - if (IsSliced(block)) { - ++sliced_blocks; - } - if (seen_blocks.contains(block)) { - continue; - } - seen_blocks.insert(block); - // We process a block and all its colocations at once below. So every - // time a block is unseen, it indicates a new colocation set. - ++colocation_sets; - - int num_sliced_colocations = (IsSliced(block) ? 1 : 0); - // GetTranstivieColocations must be called on BufferIntervals from - // buffer_intervals_. - for (const AllocationBlock* colocation : - GetTransitiveColocations(min_buffer_interval)) { - seen_blocks.insert(colocation); - if (IsSliced(colocation)) { - ++num_sliced_colocations; - } - } - if (num_sliced_colocations > 1) { - ++colocation_sets_with_multiple_sliced_blocks; - } - } - - return absl::StrCat( - "Imported repacking stats: num_blocks=", allocation_blocks_.size(), - "; num_sliced_blocks=", sliced_blocks, - "; num_colocation_sets=", colocation_sets, - "; num_colocation_sets_with_multiple_sliced_blocks=", - colocation_sets_with_multiple_sliced_blocks); - }(); - } - - BufferIntervalCompare DefaultBufferIntervalCompare() const { - return LessThanByKey([this](const BufferInterval& x) { - const BufferInterval& full_buffer_interval = - full_buffer_interval_map_.at(x.buffer); - int64_t full_buffer_interval_end = full_buffer_interval.end; - - // GetTranstivieColocations must be called on BufferIntervals from - // buffer_intervals_. - for (auto colocation : GetTransitiveColocations(x)) { - full_buffer_interval_end = - std::max(full_buffer_interval_end, - full_buffer_interval_map_.at(colocation).end); - } - - // Sort by duration (descending), size (descending), buffer (ascending). - return std::make_tuple( - full_buffer_interval.start - full_buffer_interval_end, - -full_buffer_interval.size, std::cref(*full_buffer_interval.buffer)); - }); - } - - // CommitChunks() does the following: - // 1) Commits chunks to interval_tree_. - // 2) Updates the entries in new_offsets_ and new_repacked_slicing_ for - // allocation_block, with the information from chunks. - // 3) Updates result._heap_size with the information from chunks. - // - // REQUIRED: - // - chunks is sorted in slice time order - void CommitChunks(const AllocationBlock* allocation_block, - const std::vector& chunks) { - VLOG(3) << "Committing repack chunks for " << allocation_block->ToString(); - - int64_t new_offset = -1; - std::optional repacked_slice_data = std::nullopt; - - if (IsSliced(allocation_block)) { - const SlicedAllocationData& original_slice_data = - allocation_block->original_slice_data.value(); - - // We should have a chunk per slice time. - CHECK_EQ(chunks.size(), - original_slice_data.slices_sorted_by_offset.size()); - repacked_slice_data = SlicedAllocationData(); - repacked_slice_data->slices_sorted_by_offset.reserve(chunks.size()); - - // Chunks and start times are sorted in start time order. - std::vector sorted_inclusive_start_times = - original_slice_data.SortedInclusiveStartTimes(); - for (int i = 0; i < chunks.size(); ++i) { - const Chunk& chunk = chunks[i]; - int64_t start_time = sorted_inclusive_start_times[i]; - result_.heap_size = result_.UpdatedHeapSize(chunk); - VLOG(3) << "Adding sliced chunk " << chunk.ToString() << " at [" - << start_time << ", " << allocation_block->end_time << "]"; - interval_tree_.Add(start_time, allocation_block->end_time, chunk); - new_offset = (new_offset == -1 ? chunk.offset - : std::min(new_offset, chunk.offset)); - repacked_slice_data->slices_sorted_by_offset.push_back( - AllocatedSlice({chunk.size, chunk.offset, start_time})); - } - absl::c_sort(repacked_slice_data->slices_sorted_by_offset, - [](const AllocatedSlice& lhs, const AllocatedSlice& rhs) { - return lhs.offset < rhs.offset; - }); - } else { - CHECK_EQ(chunks.size(), 1); - new_offset = chunks.front().offset; - result_.heap_size = result_.UpdatedHeapSize(chunks.front()); - VLOG(3) << "Adding unsliced chunk " << chunks.front().ToString() - << " at [" << allocation_block->inclusive_start_time << ", " - << allocation_block->end_time << ")"; - interval_tree_.Add(allocation_block->inclusive_start_time, - allocation_block->end_time, chunks.front()); - } - - CHECK_NE(new_offset, -1); - CHECK(!new_offsets_.contains(allocation_block)); - new_offsets_[allocation_block] = new_offset; - if (repacked_slice_data.has_value()) { - CHECK(IsSliced(allocation_block)); - CHECK(!new_repacked_slicing_.contains(allocation_block)); - new_repacked_slicing_[allocation_block] = *repacked_slice_data; - } - } - - // A helper structure of information we keep in FindAndCommitChunks() for - // each sliced colocation. - struct SlicedColocationData { - SlicedBufferInterval* sliced_buffer_interval; - SlicedAllocationFinder sliced_allocation_finder; - std::vector chunks; - }; - - // Finds and commits chunks for the AllocationBlock associated with - // min_buffer_interval, and all of that block's colocations. - void FindAndCommitChunks(BufferInterval* min_buffer_interval) { - const AllocationBlock* allocation_block = min_buffer_interval->buffer; - SlicedBufferInterval& sliced_buffer_interval = - sliced_buffer_interval_map_.at(allocation_block); - - int64_t max_colocation_size = GetMaxColocationSize(*min_buffer_interval); - - // Additional data we track for sliced colocations. In particular, we create - // a SlicedAllocationFinder for each sliced colocation, so we can check if - // a sliced colocation can fit at a particular offset, to accommodate the - // overestimation of free space in MakeFreeChunks(), as described at the top - // of this file. - absl::flat_hash_map - sliced_buffer_map; - for (auto colocation : - SortAllocationBlocks(GetTransitiveColocations(*min_buffer_interval))) { - if (IsSliced(colocation)) { - SlicedBufferInterval& colocation_sliced_buffer_interval = - sliced_buffer_interval_map_.at(colocation); - SlicedAllocationFinder sliced_colocation_finder = - CreateSlicedAllocationFinder( - colocation_sliced_buffer_interval, max_colocation_size, - /*preferred_offset=*/-1, - SliceTimePermutationIterator::CreateForRepack( - slice_time_permutation_iterator_type(), - GetSlicedAllocationDataPointer( - colocation->original_slice_data)), - &SlicedAllocationFinder::AllOffsetsAllowed); - sliced_buffer_map.insert(std::make_pair( - colocation, - SlicedColocationData{&colocation_sliced_buffer_interval, - std::move(sliced_colocation_finder), - /*chunks=*/{}})); - } - } - - // Using the SlicedAllocationFinders for each sliced colocation, this - // function returns true if we can place all sliced colocations at a given - // offset. - auto is_offset_allowed = [this, &sliced_buffer_map](int64_t offset) { - for (auto& block_and_colocation_data : sliced_buffer_map) { - SlicedColocationData& sliced_colocation_data = - block_and_colocation_data.second; - auto colocation_chunks = - sliced_colocation_data.sliced_allocation_finder.FindForOffset( - offset); - colocation_chunks = PostProcessFindChunkCandidatesResult( - *sliced_colocation_data.sliced_buffer_interval, - std::move(colocation_chunks)); - if (colocation_chunks.empty()) { - return false; - } - sliced_colocation_data.chunks = std::move(colocation_chunks); - } - - return true; - }; - - // Find chunks for allocation_block and its colocations. - SlicedAllocationFinder finder = CreateSlicedAllocationFinder( - sliced_buffer_interval, max_colocation_size, /*preferred_offset=*/-1, - SliceTimePermutationIterator::CreateForRepack( - slice_time_permutation_iterator_type(), - GetSlicedAllocationDataPointer( - allocation_block->original_slice_data)), - is_offset_allowed); - std::vector chunks = PostProcessFindChunkCandidatesResult( - sliced_buffer_interval, finder.Find()); - int64_t min_offset = - absl::c_min_element(chunks, [](const Chunk& lhs, const Chunk& rhs) { - return lhs.offset < rhs.offset; - })->offset; - - // Commit chunks for allocation_block. - CommitChunks(allocation_block, chunks); - - // Commit chunks for colocations. - for (auto colocation : GetTransitiveColocations(*min_buffer_interval)) { - if (IsSliced(colocation)) { - CommitChunks(colocation, sliced_buffer_map.at(colocation).chunks); - } else { - const BufferInterval& colocation_full_buffer_interval = - full_buffer_interval_map_[colocation]; - CommitChunks(colocation, - {Chunk::FromOffsetSize( - min_offset, colocation_full_buffer_interval.size)}); - } - } - } - - // We do not use result_.chunk_map, and we have our own method that combines - // finding chunks with committing them; thus, we expect this method to never - // be called. - void AddToChunkMap(const AllocationBlock* buffer, Chunk chunk) override { - LOG(FATAL) << "We should never get here."; - } - - absl::StatusOr Finish() override { - std::vector sorted_buffer_intervals = - GetSortedBufferIntervals(); - - for (auto& buffer_interval : sorted_buffer_intervals) { - if (!buffer_interval.need_allocation) { - continue; - } - - FindAndCommitChunks(&buffer_interval); - } - - Result result; - result.heap_size = result_.heap_size; - result.heap_results.push_back(result_); - return result; - } - - // A data structure for storing a chunk and its live time for use in - // debugging. - struct TimedChunk { - std::string id; - const AllocationBlock* block; - int64_t start_inclusive; - int64_t end_inclusive; - Chunk chunk; - - bool Overlaps(const TimedChunk& timed_chunk) { - if (timed_chunk.start_inclusive > end_inclusive || - timed_chunk.end_inclusive < start_inclusive) { - return false; - } - return chunk.OverlapsWith(timed_chunk.chunk); - } - }; - - void DebuggingValidate() { - std::vector timed_chunks; - for (const AllocationBlock* block : allocation_blocks_) { - if (IsSliced(block)) { - for (int i = 0; - i < block->repacked_slice_data->slices_sorted_by_offset.size(); - ++i) { - const AllocatedSlice& slice = - block->repacked_slice_data->slices_sorted_by_offset[i]; - timed_chunks.push_back( - TimedChunk{absl::StrCat(((int64_t)block), "_slice_", i), block, - slice.inclusive_start_time, block->end_time, - Chunk::FromOffsetSize(slice.offset, slice.size)}); - } - } else { - timed_chunks.push_back( - TimedChunk{absl::StrCat(((int64_t)block)), block, - block->inclusive_start_time, block->end_time, - Chunk::FromOffsetSize(block->offset, block->size)}); - } - } - - bool overlap_found = false; - for (int i = 0; i < timed_chunks.size(); ++i) { - for (int j = i + 1; j < timed_chunks.size(); ++j) { - if (timed_chunks[i].Overlaps(timed_chunks[j])) { - overlap_found = true; - LOG(ERROR) << "Allocation block overlap\n" - << " " << timed_chunks[i].block->ToString() - << "\n " << timed_chunks[j].block->ToString(); - } - } - } - - if (overlap_found) { - LOG(FATAL) << "Allocation overlap found"; - } - } - - bool Repack() { - CHECK_OK(Finish().status()); - bool success = result_.heap_size <= max_size_; - if (!success) { - VLOG(1) << "Repacking unsuccessful with heap size " << result_.heap_size; - return false; - } - - // Update AllocationBlocks. - for (AllocationBlock* block : allocation_blocks_) { - CHECK(new_offsets_.contains(block)); - block->offset = new_offsets_[block]; - if (!IsSliced(block)) { - continue; - } - - CHECK(new_repacked_slicing_.contains(block)); - block->repacked_slice_data = std::move(new_repacked_slicing_[block]); - } - - if (validate_) { - DebuggingValidate(); - } - - if (VLOG_IS_ON(2)) { - for (AllocationBlock* block : allocation_blocks_) { - VLOG(2) << "AllocationBlock after repacking: " << block->ToString(); - } - } - - VLOG(1) << "Repacking successful with heap size " << result_.heap_size; - - return true; - } - - private: - // If true, we run a potentially expensive validation to make sure there are - // no overlaps in the repacked chunks. Note, there should never be an overlap. - bool validate_ = false; - - // Maximum heap size. - int64_t max_size_; - - // Input AllocationBlocks to repack. - absl::Span allocation_blocks_; - - absl::flat_hash_map - full_buffer_interval_map_; - absl::flat_hash_map - sliced_buffer_interval_map_; - - // Data structures updated with repacking placement information as we compute - // it. - absl::flat_hash_map new_offsets_; - absl::flat_hash_map - new_repacked_slicing_; -}; - -} // namespace - -namespace memory_space_assignment { - -absl::StatusOr MemorySpaceAssignmentBestFitRepacker::Repack( - absl::Span allocations) { - BestFitRepacker best_fit_repacker = BestFitRepacker( - options_, slice_time_permutation_iterator_type_, max_size_, alignment_); - best_fit_repacker.ImportAllocationBlocks(allocations); - return best_fit_repacker.Repack(); -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.h b/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.h deleted file mode 100644 index e22daba33991f6..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.h +++ /dev/null @@ -1,77 +0,0 @@ -/* Copyright 2020 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BEST_FIT_REPACKER_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BEST_FIT_REPACKER_H_ - -#include -#include - -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/memory_space_assignment/repacking.h" - -namespace xla { -namespace memory_space_assignment { - -// This is a repacker algorithm that wraps around best fit heap algorithm in -// heap simulator. -class MemorySpaceAssignmentBestFitRepacker - : public MemorySpaceAssignmentRepacker { - public: - using BufferInterval = - GlobalDecreasingSizeBestFitHeap::BufferInterval; - using BufferIntervalCompare = - GlobalDecreasingSizeBestFitHeap::BufferIntervalCompare; - - struct BestFitRepackOptions { - // Running the validator is potentially expensive. - bool validate = false; - - // Specify the comparison function used for determining the order in which - // buffers will be allocated, during repacking. - BufferIntervalCompare buffer_interval_compare = nullptr; - }; - - MemorySpaceAssignmentBestFitRepacker( - int64_t max_size, int64_t alignment, - SliceTimePermutationIterator::Ty slice_time_permutation_iterator_type) - : MemorySpaceAssignmentRepacker(max_size, alignment), - options_(BestFitRepackOptions()), - slice_time_permutation_iterator_type_( - slice_time_permutation_iterator_type) {} - MemorySpaceAssignmentBestFitRepacker( - int64_t max_size, int64_t alignment, - SliceTimePermutationIterator::Ty slice_time_permutation_iterator_type, - BestFitRepackOptions options) - : MemorySpaceAssignmentRepacker(max_size, alignment), - options_(std::move(options)), - slice_time_permutation_iterator_type_( - slice_time_permutation_iterator_type) {} - - absl::StatusOr Repack( - absl::Span allocations) override; - - private: - BestFitRepackOptions options_; - SliceTimePermutationIterator::Ty slice_time_permutation_iterator_type_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BEST_FIT_REPACKER_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker_test.cc b/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker_test.cc deleted file mode 100644 index 45404c161a08f8..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/best_fit_repacker_test.cc +++ /dev/null @@ -1,510 +0,0 @@ -/* Copyright 2020 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/best_fit_repacker.h" - -#include -#include -#include - -#include -#include "absl/container/flat_hash_map.h" -#include "absl/types/span.h" -#include "xla/comparison_util.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "tsl/platform/test.h" - -namespace xla { - -class MemorySpaceAssignmentBestFitRepackerTest : public ::testing::Test { - protected: - MemorySpaceAssignmentBestFitRepackerTest() - : repacker_(100, 1, SliceTimePermutationIterator::Ty::kAll, options_) {} - - AllocationBlock* MakeAllocationBlock(int64_t start_time, int64_t end_time, - int64_t size, - int64_t initial_offset = -1) { - allocation_blocks_.push_back( - {start_time, end_time, size, -1, initial_offset, - static_cast(allocation_blocks_.size())}); - AllocationBlock* block = &allocation_blocks_.back(); - block->next_colocated = block; - return block; - } - - std::list allocation_blocks_; - memory_space_assignment::MemorySpaceAssignmentBestFitRepacker:: - BestFitRepackOptions options_{/*validate=*/true, - /*buffer_interval_compare=*/nullptr}; - memory_space_assignment::MemorySpaceAssignmentBestFitRepacker repacker_; -}; - -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, Simple) { - std::vector allocation_blocks; - allocation_blocks.push_back(MakeAllocationBlock(10, 20, 10)); - allocation_blocks.push_back(MakeAllocationBlock(5, 25, 15)); - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - EXPECT_EQ(allocation_blocks[0]->offset, 15); - EXPECT_EQ(allocation_blocks[1]->offset, 0); -} - -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, Colocation) { - std::vector allocation_blocks; - allocation_blocks.push_back(MakeAllocationBlock(0, 2, 10)); - allocation_blocks.push_back(MakeAllocationBlock(10, 20, 10)); - // Allocation blocks 0 and 1 are colocated. - allocation_blocks[0]->next_colocated = allocation_blocks[1]; - allocation_blocks[1]->next_colocated = allocation_blocks[0]; - allocation_blocks.push_back(MakeAllocationBlock(5, 25, 15)); - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - EXPECT_EQ(allocation_blocks[0]->offset, 15); - EXPECT_EQ(allocation_blocks[1]->offset, 15); - EXPECT_EQ(allocation_blocks[2]->offset, 0); -} - -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, TooLarge) { - // Memory size is 100, total size of buffers is 105. - std::vector allocation_blocks; - allocation_blocks.push_back(MakeAllocationBlock(10, 20, 10)); - allocation_blocks.push_back(MakeAllocationBlock(5, 25, 15)); - allocation_blocks.push_back(MakeAllocationBlock(15, 20, 10)); - allocation_blocks.push_back(MakeAllocationBlock(12, 22, 50)); - allocation_blocks.push_back(MakeAllocationBlock(10, 18, 20)); - EXPECT_FALSE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - // Make sure the buffers didn't get offset assignments. - EXPECT_EQ(allocation_blocks[0]->offset, -1); - EXPECT_EQ(allocation_blocks[1]->offset, -1); - EXPECT_EQ(allocation_blocks[2]->offset, -1); - EXPECT_EQ(allocation_blocks[3]->offset, -1); - EXPECT_EQ(allocation_blocks[4]->offset, -1); -} - -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, ColocationDifferentSizes) { - std::vector allocation_blocks; - allocation_blocks.push_back(MakeAllocationBlock(0, 2, 5)); - allocation_blocks.push_back(MakeAllocationBlock(10, 20, 10)); - // Allocation blocks 0 and 1 are colocated. - allocation_blocks[0]->next_colocated = allocation_blocks[1]; - allocation_blocks[1]->next_colocated = allocation_blocks[0]; - allocation_blocks.push_back(MakeAllocationBlock(9, 11, 2)); - allocation_blocks.push_back(MakeAllocationBlock(1, 2, 2)); - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - EXPECT_EQ(allocation_blocks[0]->offset, 0); - EXPECT_EQ(allocation_blocks[1]->offset, 0); - EXPECT_EQ(allocation_blocks[2]->offset, 10); - EXPECT_EQ(allocation_blocks[3]->offset, 5); -} - -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, RepackedSlicesFit) { - // Expected repacking: - // - // space - // ^ - // 7 | - // 6 | +-----+ - // 5 | +-------+-++ E| - // 4 | | B |+--++--++--+ - // 3 | | || || F|| | - // 2 +----------+---++----++ |+--++ | - // 1 | A || C || D | - // 0 +--------------++--------++------+ - // +----|----|----|----|----|----|----|--> time - // 0 5 10 15 20 25 30 35 - - std::vector allocation_blocks; - // Block A - allocation_blocks.push_back(MakeAllocationBlock(0, 15, 2)); - // Block B - allocation_blocks.push_back(MakeAllocationBlock(11, 21, 3)); - // Block C - allocation_blocks.push_back(MakeAllocationBlock(16, 25, 4)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, -1, 16}, AllocatedSlice{2, -1, 22}}}); - // Block D - allocation_blocks.push_back(MakeAllocationBlock(26, 33, 4)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, -1, 26}, AllocatedSlice{2, -1, 30}}}); - // Block E - allocation_blocks.push_back(MakeAllocationBlock(19, 25, 2)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{1, -1, 19}, AllocatedSlice{1, -1, 22}}}); - // Block F - allocation_blocks.push_back(MakeAllocationBlock(26, 29, 2)); - - // Specify the repacking sort order as the order in which blocks were added to - // allocation_blocks. - // - By placing C after B, we test that the repacker can place a sliced block - // around another block - // - By placing F after D, we test that the repacker can fill in the extra - // space left behind by slicing. - absl::flat_hash_map sort_keys; - for (int i = 0; i < allocation_blocks.size(); ++i) { - sort_keys[allocation_blocks[i]] = i; - } - options_.buffer_interval_compare = LessThanByKey( - [sort_keys](const memory_space_assignment:: - MemorySpaceAssignmentBestFitRepacker::BufferInterval& x) { - return sort_keys.at(x.buffer); - }); - repacker_ = memory_space_assignment::MemorySpaceAssignmentBestFitRepacker( - 100, 1, SliceTimePermutationIterator::Ty::kAll, options_); - - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - // Check results - // Block A - EXPECT_EQ(allocation_blocks[0]->offset, 0); - EXPECT_FALSE(allocation_blocks[0]->repacked_slice_data.has_value()); - // Block B - EXPECT_EQ(allocation_blocks[1]->offset, 2); - EXPECT_FALSE(allocation_blocks[1]->repacked_slice_data.has_value()); - // Block C - EXPECT_EQ(allocation_blocks[2]->offset, 0); - ASSERT_TRUE(allocation_blocks[2]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[2]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{2, 0, 16}, AllocatedSlice{2, 2, 22}}}))); - // Block D - EXPECT_EQ(allocation_blocks[3]->offset, 0); - ASSERT_TRUE(allocation_blocks[3]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[3]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{2, 0, 26}, AllocatedSlice{2, 2, 30}}}))); - // Block E - EXPECT_EQ(allocation_blocks[4]->offset, 4); - ASSERT_TRUE(allocation_blocks[4]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[4]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{1, 4, 22}, AllocatedSlice{1, 5, 19}}}))); - // Block F - EXPECT_EQ(allocation_blocks[5]->offset, 2); - EXPECT_FALSE(allocation_blocks[5]->repacked_slice_data.has_value()); -} - -// Test that we do not permute slice start times in a way that changes the -// original slice size-start time mappings. Doing so breaks assumptions that -// MSA uses to construct its internal state prior to repacking. -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, - SliceTimePermutationsMatchOriginalSizeTimeMapping) { - // Original placement: Ideal repacking, but unsupported: - // - // space space - // ^ ^ - // 7 | +---------+ 7 | - // 6 | | C | 6 | +---------+ - // 5 | +-----+---+ 5 | | C | - // 4 | +-----+ | 4 | +---------+ - // 3 | | B | 3 | | B | - // 2 +----+----+----+ 2 +----+----++ | - // 1 | A | 1 | A |+---+ - // 0 +---------+ 0 +---------+ - // +----|----|----|----> time +----|----|----|----> time - // 0 5 10 15 0 5 10 15 - - std::vector allocation_blocks; - // Block A - allocation_blocks.push_back(MakeAllocationBlock(0, 10, 2, 0)); - // Block B - allocation_blocks.push_back(MakeAllocationBlock(5, 15, 3, 2)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, 2, 5}, AllocatedSlice{1, 4, 11}}}); - // Block C - allocation_blocks.push_back(MakeAllocationBlock(5, 15, 2, 6)); - - // Specify the repacking sort order as the order in which blocks were added to - // allocation_blocks. We need to do this so that B is placed before C. If C - // is placed before B, C will sit directly on top of A, and the repacker would - // never try to permute B's slice size-start time mapping. - absl::flat_hash_map sort_keys; - for (int i = 0; i < allocation_blocks.size(); ++i) { - sort_keys[allocation_blocks[i]] = i; - } - options_.buffer_interval_compare = LessThanByKey( - [sort_keys](const memory_space_assignment:: - MemorySpaceAssignmentBestFitRepacker::BufferInterval& x) { - return sort_keys.at(x.buffer); - }); - repacker_ = memory_space_assignment::MemorySpaceAssignmentBestFitRepacker( - 100, 1, SliceTimePermutationIterator::Ty::kAll, options_); - - // The repacker returns true as long as the result fits in the max size, - // regardless of whether it has actually changed anything. - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - // Typically the heap_simulator would prefer to start Block B at a smaller - // offset, i.e., offset 1 rather than offset 2. However, in order to do so, - // the repacker would have to permute the original slice size-start time - // mapping, which is not permitted. Thus, we ensure that the repacked B's - // larger slice is assigned the smaller offset and earlier start time. - ASSERT_TRUE(allocation_blocks[1]->repacked_slice_data.has_value()); - ASSERT_EQ( - allocation_blocks[1]->repacked_slice_data->slices_sorted_by_offset.size(), - 2); - const AllocatedSlice& slice_with_smaller_offset = - allocation_blocks[1]->repacked_slice_data->slices_sorted_by_offset[0]; - const AllocatedSlice& slice_with_larger_offset = - allocation_blocks[1]->repacked_slice_data->slices_sorted_by_offset[1]; - // The larger slice is assigned to the smaller offset. - ASSERT_GT(slice_with_smaller_offset.size, slice_with_larger_offset.size); - const AllocatedSlice& larger_slice = slice_with_smaller_offset; - const AllocatedSlice& smaller_slice = slice_with_larger_offset; - // The larger slice is assigned to the earlier start time. - ASSERT_LT(larger_slice.inclusive_start_time, - smaller_slice.inclusive_start_time); -} - -// Test that we do not permute slice start times in a way that changes the -// original slice size-start time mappings. Doing so breaks assumptions that -// MSA uses to construct its internal state prior to repacking. -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, - SliceTimePermutationsMatchOriginalSizeTimeMapping2) { - // Original placement: New placement: - // - // space space - // ^ ^ - // 7 | 7 | - // 6 | +--------+ 6 | - // 5 | | B | 5 | +---------+ - // 4 | +-----+---+----+ 4 | | C | - // 3 | | C | 3 | +-----+ | - // 2 +----+----++ | 2 +---------++---+----+ - // 1 | A |+---+ 1 | A || B | - // 0 +---------+ 0 +---------++--------+ - // +----|----|----|----|--> time +----|----|----|----|--> time - // 0 5 10 15 20 0 5 10 15 20 - - std::vector allocation_blocks; - // Block A - allocation_blocks.push_back(MakeAllocationBlock(0, 10, 2, 0)); - // Block B - allocation_blocks.push_back(MakeAllocationBlock(11, 20, 2, 4)); - // Block C - allocation_blocks.push_back(MakeAllocationBlock(5, 15, 3, 1)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{1, 1, 5}, AllocatedSlice{2, 2, 11}}}); - - // Specify the repacking sort order as the order in which blocks were added to - // allocation_blocks. We need to do this so that B is placed before C. - absl::flat_hash_map sort_keys; - for (int i = 0; i < allocation_blocks.size(); ++i) { - sort_keys[allocation_blocks[i]] = i; - } - options_.buffer_interval_compare = LessThanByKey( - [sort_keys](const memory_space_assignment:: - MemorySpaceAssignmentBestFitRepacker::BufferInterval& x) { - return sort_keys.at(x.buffer); - }); - repacker_ = memory_space_assignment::MemorySpaceAssignmentBestFitRepacker( - 100, 1, SliceTimePermutationIterator::Ty::kAll, options_); - - // The repacker returns true as long as the result fits in the max size, - // regardless of whether it has actually changed anything. - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - // Check results - // - // Typically the heap_simulator would prefer to start the first slice of - // Block C at time 5 and the second block at time 11, but that is not allowed - // because it permutes the original slice size-time mapping. - // - // Block A - EXPECT_EQ(allocation_blocks[0]->offset, 0); - EXPECT_FALSE(allocation_blocks[0]->repacked_slice_data.has_value()); - // Block B - EXPECT_EQ(allocation_blocks[1]->offset, 0); - EXPECT_FALSE(allocation_blocks[1]->repacked_slice_data.has_value()); - // Block C - EXPECT_EQ(allocation_blocks[2]->offset, 2); - ASSERT_TRUE(allocation_blocks[2]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[2]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{1, 2, 5}, AllocatedSlice{2, 3, 11}}}))); -} - -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, SlicedColocationsFit) { - // Expected repacking: - // - // space - // ^ - // 9 | +-+ - // 8 | | | - // 7 | |F|+-+ - // 6 | +-----++-+| || | - // 5 | | C || || || | - // 4 +----+--++-++ |+-++ | - // 3 | B || E || D | - // 2 +-------++--+-++----+ - // 1 | A | - // 0 +-----------+ - // +----|----|----|----|----|----|----|--> time - // 0 5 10 15 20 25 30 35 - - // D is colocated with E; thus, they will be allocated together. - std::vector allocation_blocks; - // Block A - allocation_blocks.push_back(MakeAllocationBlock(0, 12, 2)); - // Block B - allocation_blocks.push_back(MakeAllocationBlock(0, 8, 2)); - // Block C - allocation_blocks.push_back(MakeAllocationBlock(5, 11, 2)); - // Block D - allocation_blocks.push_back(MakeAllocationBlock(15, 20, 5)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, -1, 15}, AllocatedSlice{3, -1, 18}}}); - // Block E - allocation_blocks.push_back(MakeAllocationBlock(9, 14, 4)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, -1, 9}, AllocatedSlice{2, -1, 12}}}); - // Colocate E with D. - allocation_blocks.back()->next_colocated = allocation_blocks[3]; - allocation_blocks[3]->next_colocated = allocation_blocks.back(); - // Block F - allocation_blocks.push_back(MakeAllocationBlock(15, 17, 5)); - - // Specify the repacking sort order as the order in which blocks were added to - // allocation_blocks. - // - By placing E after C, we test that the repacker can place a sliced block - // around another block - // - By placing F after D, we test that the repacker can fill in the extra - // space left behind by slicing. - absl::flat_hash_map sort_keys; - for (int i = 0; i < allocation_blocks.size(); ++i) { - sort_keys[allocation_blocks[i]] = i; - } - options_.buffer_interval_compare = LessThanByKey( - [sort_keys](const memory_space_assignment:: - MemorySpaceAssignmentBestFitRepacker::BufferInterval& x) { - return sort_keys.at(x.buffer); - }); - repacker_ = memory_space_assignment::MemorySpaceAssignmentBestFitRepacker( - 100, 1, SliceTimePermutationIterator::Ty::kAll, options_); - - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - // Check results - // Block A - EXPECT_EQ(allocation_blocks[0]->offset, 0); - EXPECT_FALSE(allocation_blocks[0]->repacked_slice_data.has_value()); - // Block B - EXPECT_EQ(allocation_blocks[1]->offset, 2); - EXPECT_FALSE(allocation_blocks[1]->repacked_slice_data.has_value()); - // Block C - EXPECT_EQ(allocation_blocks[2]->offset, 4); - ASSERT_FALSE(allocation_blocks[2]->repacked_slice_data.has_value()); - // Block D - EXPECT_EQ(allocation_blocks[3]->offset, 2); - ASSERT_TRUE(allocation_blocks[3]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[3]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{2, 2, 15}, AllocatedSlice{3, 4, 18}}}))); - // Block E - EXPECT_EQ(allocation_blocks[4]->offset, 2); - ASSERT_TRUE(allocation_blocks[4]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[4]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{2, 2, 9}, AllocatedSlice{2, 4, 12}}}))); - // Block F - EXPECT_EQ(allocation_blocks[5]->offset, 4); - EXPECT_FALSE(allocation_blocks[5]->repacked_slice_data.has_value()); -} - -// Test that we do not permute slice start times in a way that changes the -// original slice size-start time mappings. Doing so breaks assumptions that -// MSA uses to construct its internal state prior to repacking. -TEST_F(MemorySpaceAssignmentBestFitRepackerTest, - SlicedColocationsPermutationsMatchOriginalSizeTimeMapping) { - // Original placement: Ideal repacking, but unsupported: - // - // space space - // ^ ^ - // 8 | 8 | - // 7 |+--------+ +---+ 7 | - // 6 || | | | 6 | - // 5 || C | | D | 5 |+--------++--------+ - // 4 |+----+ |+----+ | 4 || || | - // 3 | | || | 3 || C || D | - // 2 |+---++---++---+----+ 2 |+---++ |+---++ | - // 1 || A | | B | 1 || A || || B || | - // 0 |+---+ +---+ 0 |+---++---++---++---+ - // +----|----|----|----|----> time +----|----|----|----|----> time - // 0 5 10 15 20 0 5 10 15 20 - - std::vector allocation_blocks; - // Block A - allocation_blocks.push_back(MakeAllocationBlock(1, 5, 2)); - // Block B - allocation_blocks.push_back(MakeAllocationBlock(11, 15, 2)); - // Block C - allocation_blocks.push_back(MakeAllocationBlock(1, 10, 5)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, 2, 6}, AllocatedSlice{3, 4, 1}}}); - // Block D - allocation_blocks.push_back(MakeAllocationBlock(15, 20, 5)); - allocation_blocks.back()->original_slice_data = SlicedAllocationData( - {{AllocatedSlice{2, 2, 11}, AllocatedSlice{3, 4, 16}}}); - // Colocate D with C. - allocation_blocks.back()->next_colocated = allocation_blocks[2]; - allocation_blocks[2]->next_colocated = allocation_blocks.back(); - - // Specify the repacking sort order as the order in which blocks were added to - // allocation_blocks. By placing A and B before C/D, the repacker will try - // permutations of C/D's slices that fit around A and B. - absl::flat_hash_map sort_keys; - for (int i = 0; i < allocation_blocks.size(); ++i) { - sort_keys[allocation_blocks[i]] = i; - } - options_.buffer_interval_compare = LessThanByKey( - [sort_keys](const memory_space_assignment:: - MemorySpaceAssignmentBestFitRepacker::BufferInterval& x) { - return sort_keys.at(x.buffer); - }); - repacker_ = memory_space_assignment::MemorySpaceAssignmentBestFitRepacker( - 100, 1, SliceTimePermutationIterator::Ty::kAll, options_); - - EXPECT_TRUE(*repacker_.Repack(absl::MakeSpan(allocation_blocks))); - - // Check results - // - // Typically the heap simulator would like start C/D at offset 0, which is - // lower than C/D's actual placement at offset 2. However, in order to place - // C/D at offset 0, we would need to permute the slice-time mappings of the - // colocation D, which is not permitted. - // - // Block A - EXPECT_EQ(allocation_blocks[0]->offset, 0); - EXPECT_FALSE(allocation_blocks[0]->repacked_slice_data.has_value()); - // Block B - EXPECT_EQ(allocation_blocks[1]->offset, 0); - EXPECT_FALSE(allocation_blocks[1]->repacked_slice_data.has_value()); - // Block C - EXPECT_EQ(allocation_blocks[2]->offset, 2); - ASSERT_TRUE(allocation_blocks[2]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[3]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{2, 2, 11}, AllocatedSlice{3, 4, 16}}}))); - // Block D - EXPECT_EQ(allocation_blocks[3]->offset, 2); - ASSERT_TRUE(allocation_blocks[3]->repacked_slice_data.has_value()); - EXPECT_EQ(*allocation_blocks[3]->repacked_slice_data, - (SlicedAllocationData( - {{AllocatedSlice{2, 2, 11}, AllocatedSlice{3, 4, 16}}}))); -} - -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.cc b/third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.cc deleted file mode 100644 index 21a0f736dff36a..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.cc +++ /dev/null @@ -1,172 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/buffer_interval_comparator.h" - -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" -#include "absl/strings/string_view.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/utils.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -MemoryBoundednessBufferIntervalComparator:: - MemoryBoundednessBufferIntervalComparator( - const CostAnalysis& cost_analysis, - CostAnalysis::Cache* cost_analysis_cache) - : BufferIntervalComparator(), - cost_analysis_(cost_analysis), - cost_analysis_cache_(cost_analysis_cache) {} - -MemoryBoundednessBufferIntervalComparator:: - MemoryBoundednessBufferIntervalComparator( - const CostAnalysis& cost_analysis, - CostAnalysis::Cache* cost_analysis_cache, - MsaSortOrderOverrides msa_sort_order_overrides) - : BufferIntervalComparator(), - cost_analysis_(cost_analysis), - cost_analysis_cache_(cost_analysis_cache), - msa_sort_order_overrides_(msa_sort_order_overrides) {} - -std::string -MemoryBoundednessBufferIntervalComparator::DescribeComparisonCriteria() const { - return "[override priority, -memory boundedness, -size, -buffer duration, " - "latest use time, (inclusive) start time, instruction id ]"; -} - -std::string MemoryBoundednessBufferIntervalComparator::CriteriaToString( - const MsaBufferInterval& buffer_interval) { - return absl::StrCat("[ ", absl::StrJoin(GetTuple(buffer_interval), ", "), - " ]"); -} - -bool MemoryBoundednessBufferIntervalComparator::LessThan( - const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - return GetTuple(lhs) < GetTuple(rhs); -} - -int64_t MemoryBoundednessBufferIntervalComparator::GetLatestUseTime( - const MsaBufferInterval& buffer_interval) { - auto latest_use_it = buffer_to_latest_use_.find(buffer_interval.buffer); - if (latest_use_it == buffer_to_latest_use_.end()) { - int64_t latest_use_time = 0; - for (const HloUse& use : buffer_interval.buffer->GetUses()) { - auto it = cost_analysis_.hlo_live_range().instruction_schedule().find( - use.instruction); - if (it != cost_analysis_.hlo_live_range().instruction_schedule().end()) { - latest_use_time = std::max(latest_use_time, it->second); - } - } - latest_use_it = - buffer_to_latest_use_ - .insert(std::make_pair(buffer_interval.buffer, latest_use_time)) - .first; - } - return latest_use_it->second; -} - -MemoryBoundednessBufferIntervalComparator::ComparisonTuple -MemoryBoundednessBufferIntervalComparator::GetTuple( - const MsaBufferInterval& buffer_interval) { - int64_t priority = - MemorySpaceAssignmentUtils::GetBufferIntervalOverridePriority( - msa_sort_order_overrides_, buffer_interval); - float inverse_memory_boundedness = - -1.0 * cost_analysis_.GetMemoryBoundedness(buffer_interval, - cost_analysis_cache_); - int64_t inverse_buffer_size = -1 * buffer_interval.size; - int64_t inverse_buffer_duration = buffer_interval.start - buffer_interval.end; - int64_t latest_use_time = GetLatestUseTime(buffer_interval); - int64_t buffer_start_time = buffer_interval.start; - auto buffer_id = buffer_interval.buffer->id(); - return std::make_tuple(priority, inverse_memory_boundedness, - inverse_buffer_size, inverse_buffer_duration, - latest_use_time, buffer_start_time, buffer_id); -} - -DefaultCrossProgramPrefetchBufferIntervalComparator:: - DefaultCrossProgramPrefetchBufferIntervalComparator( - const HloLiveRange& hlo_live_range, - const MsaSortOrderOverrides& msa_sort_order_overrides) - : BufferIntervalComparator(), - hlo_live_range_(hlo_live_range), - msa_sort_order_overrides_(msa_sort_order_overrides) {} - -std::string DefaultCrossProgramPrefetchBufferIntervalComparator:: - DescribeComparisonCriteria() const { - return "[ -size, -cumulative use size, latest use, instruction id]"; -} - -std::string -DefaultCrossProgramPrefetchBufferIntervalComparator::CriteriaToString( - const MsaBufferInterval& buffer_interval) { - return absl::StrCat("[ ", absl::StrJoin(GetTuple(buffer_interval), ", "), - " ]"); -} - -bool DefaultCrossProgramPrefetchBufferIntervalComparator::LessThan( - const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - return GetTuple(lhs) < GetTuple(rhs); -} - -DefaultCrossProgramPrefetchBufferIntervalComparator::ComparisonTuple -DefaultCrossProgramPrefetchBufferIntervalComparator::GetTuple( - const MsaBufferInterval& buffer_interval) { - int64_t priority = - MemorySpaceAssignmentUtils::GetBufferIntervalOverridePriority( - msa_sort_order_overrides_, buffer_interval, - /*is_cross_program_prefetch=*/true); - auto sort_data_it = additional_sort_data_.find(buffer_interval.buffer); - if (sort_data_it == additional_sort_data_.end()) { - AdditionalSortData sort_data; - absl::c_for_each(buffer_interval.buffer->GetUses(), [&](const HloUse& use) { - auto it = hlo_live_range_.instruction_schedule().find(use.instruction); - if (it == hlo_live_range_.instruction_schedule().end()) { - return; - } - sort_data.latest_use = std::max(sort_data.latest_use, it->second); - sort_data.cumulative_use_size += - ShapeUtil::ElementsInRecursive(use.instruction->shape()); - }); - sort_data_it = - additional_sort_data_.try_emplace(buffer_interval.buffer, sort_data) - .first; - } - - return std::make_tuple(priority, -1 * buffer_interval.size, - -1 * sort_data_it->second.cumulative_use_size, - sort_data_it->second.latest_use, - buffer_interval.buffer->id()); -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.h b/third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.h deleted file mode 100644 index 5c7a94b6ffd468..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.h +++ /dev/null @@ -1,150 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BUFFER_INTERVAL_COMPARATOR_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BUFFER_INTERVAL_COMPARATOR_H_ - -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/buffer_value.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/utils.h" - -namespace xla { -namespace memory_space_assignment { - -using MsaBufferIntervalCompare = - GlobalDecreasingSizeBestFitHeap::BufferIntervalCompare; - -// The MsaBufferInterval sorting interface that MemorySpaceAssignment expects. -class BufferIntervalComparator { - public: - virtual ~BufferIntervalComparator() = default; - - // A logging string explaining the sorting criteria. E.g., [ -size, offset ] - // indicates we sort (desc) size, then (asc) offset. - virtual std::string DescribeComparisonCriteria() const = 0; - - // A logging string containing the values used to sort buffer_interval. - // E.g., we might return [ -1024, 100 ], if the criteria is [ -size, - // offset ]. - virtual std::string CriteriaToString( - const MsaBufferInterval& buffer_interval) = 0; - - // comparator.LessThan(lhs, rhs) will be used for MsaBufferIntervalCompare. - virtual bool LessThan(const MsaBufferInterval& lhs, - const MsaBufferInterval& rhs) = 0; - - // Used to create a functor that can be passed to a method like std::sort. - // E.g., absl::c_sort(v, comparator.GetComparisonFunctor()); - MsaBufferIntervalCompare GetComparisonFunctor() { - return [this](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - return LessThan(lhs, rhs); - }; - } - - protected: - BufferIntervalComparator() = default; -}; - -// A BufferIntervalComparator that utilizes MemoryBoundedness as its primary -// sorting criteria. -// -// This comparator caches HloValues -> latest use time. -class MemoryBoundednessBufferIntervalComparator - : public BufferIntervalComparator { - public: - MemoryBoundednessBufferIntervalComparator( - const CostAnalysis& cost_analysis, - CostAnalysis::Cache* cost_analysis_cache); - - MemoryBoundednessBufferIntervalComparator( - const CostAnalysis& cost_analysis, - CostAnalysis::Cache* cost_analysis_cache, - MsaSortOrderOverrides msa_sort_order_overrides); - - ~MemoryBoundednessBufferIntervalComparator() override = default; - - std::string DescribeComparisonCriteria() const override; - std::string CriteriaToString( - const MsaBufferInterval& buffer_interval) override; - bool LessThan(const MsaBufferInterval& lhs, - const MsaBufferInterval& rhs) override; - - private: - // See the value returned by DescribeComparisonCriteria() for the meaning of - // each tuple element. - using ComparisonTuple = std::tuple; - - ComparisonTuple GetTuple(const MsaBufferInterval& buffer_interval); - int64_t GetLatestUseTime(const MsaBufferInterval& buffer_interval); - absl::flat_hash_map buffer_to_latest_use_; - const CostAnalysis& cost_analysis_; - CostAnalysis::Cache* cost_analysis_cache_; - - // Config to override alternate memory assignment sorting order for filtered - // buffers. - MsaSortOrderOverrides msa_sort_order_overrides_; -}; - -// The default BufferIntervalComparator used for cross-program prefetching. -// -// This class caches HloValue -> {latest use, cumulative use size }. -class DefaultCrossProgramPrefetchBufferIntervalComparator - : public BufferIntervalComparator { - public: - explicit DefaultCrossProgramPrefetchBufferIntervalComparator( - const HloLiveRange& hlo_live_range, - const MsaSortOrderOverrides& msa_sort_order_overrides); - - ~DefaultCrossProgramPrefetchBufferIntervalComparator() override = default; - - std::string DescribeComparisonCriteria() const override; - std::string CriteriaToString( - const MsaBufferInterval& buffer_interval) override; - bool LessThan(const MsaBufferInterval& lhs, - const MsaBufferInterval& rhs) override; - - private: - // See the value returned by DescribeComparisonCriteria() for the meaning of - // each tuple element. - using ComparisonTuple = - std::tuple; - - struct AdditionalSortData { - int64_t latest_use = 0; - int64_t cumulative_use_size = 0; - }; - - ComparisonTuple GetTuple(const MsaBufferInterval& buffer_interval); - - absl::flat_hash_map - additional_sort_data_; - const HloLiveRange& hlo_live_range_; - const MsaSortOrderOverrides& msa_sort_order_overrides_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_BUFFER_INTERVAL_COMPARATOR_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/cost_analysis.cc b/third_party/xla/xla/service/memory_space_assignment/cost_analysis.cc deleted file mode 100644 index 7543f1fdca0450..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/cost_analysis.cc +++ /dev/null @@ -1,502 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/cost_analysis.h" - -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_set.h" -#include "absl/log/check.h" -#include "absl/memory/memory.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/analysis/while_loop_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/utils/hlo_live_range.h" -#include "xla/service/call_graph.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/utils.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -/*static*/ absl::StatusOr> CostAnalysis::Create( - OpCostManager& op_cost_manager, const CostAnalysisOptions& options, - const AliasInfo* alias_info, const HloModule& module, - HloAliasAnalysis* alias_analysis) { - ABSL_ASSIGN_OR_RETURN(auto hlo_live_range, - HloLiveRange::Run(module.schedule(), *alias_analysis, - module.entry_computation())); - auto call_graph = CallGraph::Build(&module); - // Using `new` to access a non-public constructor. - return absl::WrapUnique( - new CostAnalysis(op_cost_manager, options, alias_analysis, - std::move(hlo_live_range), std::move(call_graph))); -} - -CostAnalysis::CostAnalysis(OpCostManager& op_cost_manager, - const CostAnalysisOptions& options, - HloAliasAnalysis* alias_analysis, - std::unique_ptr hlo_live_range, - std::unique_ptr call_graph) - : op_cost_manager_(op_cost_manager), - options_(options), - alias_analysis_(alias_analysis), - hlo_live_range_(std::move(hlo_live_range)), - call_graph_(std::move(call_graph)) {} - -int64_t CostAnalysis::GetShapeSizeBytes(const Shape& shape) const { - return options_.shape_size_bytes_fn(shape); -} - -double CostAnalysis::DefaultMemBandwidthBytesPerSecond( - bool use_scaling_factor) const { - if (use_scaling_factor) { - return options_.async_copy_bandwidth_scaling_factor * - options_.default_mem_bandwidth_bytes_per_second; - } - return options_.default_mem_bandwidth_bytes_per_second; -} - -float CostAnalysis::OperandBytesAccessed(const HloInstruction& instruction, - int64_t operand_num, - const ShapeIndex& shape_index) const { - return op_cost_manager_.OperandBytesAccessed(instruction, operand_num, - shape_index); -} - -float CostAnalysis::OutputBytesAccessed(const HloInstruction& instruction, - const ShapeIndex& shape_index) const { - return op_cost_manager_.OutputBytesAccessed(instruction, shape_index); -} - -float CostAnalysis::GetAlternateMemoryBenefit( - const HloInstruction& instruction, float elapsed_time_due_to_alternate_mem, - CostAnalysis::Cache* cache) const { - float elapsed_time_due_to_compute = - GetInstructionElapsedDueToCompute(instruction); - float elapsed_time_due_to_memory = - GetInstructionElapsedDueToMemory(instruction); - if (elapsed_time_due_to_memory > elapsed_time_due_to_compute) { - // Memory bound, return how much alternate memory is better. - float while_nest_multiplier; - if (cache) { - // If there is a cache provided, memoize the while nest multiplier. - auto it = cache->while_nest_multiplier.find(&instruction); - if (it != cache->while_nest_multiplier.end()) { - while_nest_multiplier = it->second; - } else { - while_nest_multiplier = GetWhileNestMultiplier( - CalculateComputationNestLevel(&instruction, - /*while_only=*/true)); - cache->while_nest_multiplier[&instruction] = while_nest_multiplier; - } - } else { - while_nest_multiplier = GetWhileNestMultiplier( - CalculateComputationNestLevel(&instruction, - /*while_only=*/true)); - } - return (elapsed_time_due_to_memory - elapsed_time_due_to_alternate_mem) * - while_nest_multiplier; - } else { - // Compute bound, return how far off are we to memory boundedness. - return elapsed_time_due_to_memory - elapsed_time_due_to_compute; - } -} - -float CostAnalysis::GetMemoryBoundedness( - const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval, - CostAnalysis::Cache* cache) const { - if (cache) { - auto it = - cache->memory_boundedness.find(interval.buffer->defining_position()); - if (it != cache->memory_boundedness.end()) { - return it->second; - } - } - float alternate_mem_benefit = - GetAlternateMemoryBenefit(interval.buffer->defining_position(), cache); - - for (const HloBuffer* buffer : alias_analysis_->ComputeBuffersAt( - interval.buffer->defining_position().instruction, - interval.buffer->defining_position().index)) { - for (const HloValue* value : buffer->values()) { - for (const HloUse& use : value->GetUses()) { - // We look inside the called computations of while and conditional, so - // don't use the benefit of while and conditional directly. - if (use.instruction->opcode() == HloOpcode::kWhile || - use.instruction->opcode() == HloOpcode::kConditional) { - continue; - } - float use_alternate_mem_benefit = GetAlternateMemoryBenefit(use, cache); - // If the benefit is positive (memory bound), add it to this buffer's - // benefit. If the benefit is negative (compute bound), calculate the - // maximum. - if (alternate_mem_benefit > 0 && use_alternate_mem_benefit > 0) { - alternate_mem_benefit += use_alternate_mem_benefit; - } else { - alternate_mem_benefit = - std::max(alternate_mem_benefit, use_alternate_mem_benefit); - } - } - } - } - - // Penalize larger buffers by dividing the benefit by the square root of - // the size. Empirically, we observed this resulted in better performance - // compared to dividing by the size. - float memory_boundedness = 1; - if (options_ - .xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers == - "NO_SCALE") { - memory_boundedness = alternate_mem_benefit; - } else { - memory_boundedness = alternate_mem_benefit / std::sqrt(interval.size); - } - - if (cache) { - cache->memory_boundedness[interval.buffer->defining_position()] = - memory_boundedness; - } - return memory_boundedness; -} - -float CostAnalysis::GetAlternateMemoryBenefit( - const HloPosition& position, CostAnalysis::Cache* cache) const { - return GetAlternateMemoryBenefit( - *position.instruction, - GetInstructionElapsedDueToMemory( - *position.instruction, - /*operands_in_alternate_mem=*/{}, - /*outputs_in_alternate_mem=*/{position.index}), - cache); -} - -float CostAnalysis::GetAlternateMemoryBenefit( - const HloUse& use, CostAnalysis::Cache* cache) const { - return GetAlternateMemoryBenefit( - *use.instruction, - GetInstructionElapsedDueToMemory( - *use.instruction, - /*operands_in_alternate_mem=*/{std::make_pair(use.operand_number, - use.operand_index)}), - cache); -} - -int CostAnalysis::CalculateComputationNestLevel( - const HloInstruction* instruction, bool while_only) const { - int nest_level = 0; - const HloComputation* computation = instruction->parent(); - while (!computation->IsEntryComputation()) { - auto& node = call_graph_->GetNode(computation); - auto callsites = node.caller_callsites(); - CHECK(node.computation()->IsAsyncComputation() || callsites.size() == 1) - << "The module is not flattened!"; - auto& callsite = callsites[0]; - if (!while_only || callsite.instruction()->opcode() == HloOpcode::kWhile) { - ++nest_level; - } - computation = callsite.instruction()->parent(); - } - return nest_level; -} - -// TODO(hanruobing): This function assumes all nested layers have the -// same hard-coded trip count for simplicity. I plan to replace it with the -// more accurate function (CalculateNestTripCount). -float CostAnalysis::GetWhileNestMultiplier(int while_nest_level) const { - return IPow( - options_.xla_tpu_memory_space_assignment_while_execution_count, - while_nest_level); -} - -float CostAnalysis::CalculateNestTripCount(const HloInstruction* instruction, - CostAnalysis::Cache* cache) const { - float total_trip_count = 1.0; - const HloComputation* computation = instruction->parent(); - while (!computation->IsEntryComputation()) { - if (cache) { - auto it = cache->computation_trip_count.find(computation); - if (it != cache->computation_trip_count.end()) { - if (computation == instruction->parent()) { - return it->second; - } else { - total_trip_count *= it->second; - break; - } - } - } - CallGraphNode& node = call_graph_->GetNode(computation); - absl::Span callsites = node.caller_callsites(); - const xla::CallSite& callsite = callsites[0]; - if (callsite.instruction()->opcode() == HloOpcode::kWhile) { - HloInstruction* while_op = callsite.instruction(); - std::optional trip_count; - if (!trip_count.has_value()) { - // TODO(hanruobing): Apply PrepareModuleForUnrolling on the module may - // provide more accurate results for trip count analysis . However, it - // may downgrade the performance of MSA. We need more evaluation to - // decide whether to apply that pass before MSA. - trip_count = ComputeWhileLoopTripCount(while_op); - } - total_trip_count *= trip_count.value_or( - options_.xla_tpu_memory_space_assignment_while_execution_count); - } - computation = callsite.instruction()->parent(); - } - if (cache) { - cache->computation_trip_count[instruction->parent()] = total_trip_count; - } - return total_trip_count; -} - -float CostAnalysis::GetDefaultMemoryAccessOverhead( - const HloInstruction& instruction, - absl::Span> operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const { - // Calculate the pipeline overhead of accessing the default memory. We use the - // maximum of the window size heuristic and the actual default memory bytes - // accessed multiplied with the compute as the overhead. So, the math is: - // - // overhead = compute_per_iteration - // = compute_elapsed / num_iterations - // = compute_elapsed / (bytes_accessed / window_size) - // = (window_size / bytes_accessed) * compute_elapsed - const float window_size_bytes = - options_.pipeline_overhead_window_size_mib * 1024 * 1024; - const float bytes_accessed = op_cost_manager_.TotalBytesAccessed(instruction); - const float default_memory_bytes_accessed = - bytes_accessed - - GetBytesAccessedFromAlternateMemory( - instruction, operands_in_alternate_mem, outputs_in_alternate_mem); - const float compute_elapsed = GetInstructionElapsedDueToCompute(instruction); - const float effective_window_size_bytes = - std::min(window_size_bytes, default_memory_bytes_accessed); - float overhead = 0; - if (bytes_accessed > 0) { - overhead = (effective_window_size_bytes / bytes_accessed) * compute_elapsed; - } - return overhead; -} - -float CostAnalysis::GetDefaultMemoryBandwidthIdleTime( - const HloInstruction& instruction, - absl::Span> operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const { - const float default_memory_bytes_accessed = - op_cost_manager_.TotalBytesAccessed(instruction) - - GetBytesAccessedFromAlternateMemory( - instruction, operands_in_alternate_mem, outputs_in_alternate_mem); - const float elapsed_due_to_default_mem = - default_memory_bytes_accessed / DefaultMemBandwidthBytesPerSecond(); - const float elapsed = GetInstructionElapsedInAlternateMemory( - instruction, operands_in_alternate_mem, outputs_in_alternate_mem); - return elapsed - elapsed_due_to_default_mem; -} - -float CostAnalysis::GetBytesAccessedFromAlternateMemory( - const HloInstruction& instruction, - absl::Span> operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const { - float bytes_accessed_from_alternate_mem = 0.0; - for (auto& operand : operands_in_alternate_mem) { - const float operand_bytes_accessed = op_cost_manager_.OperandBytesAccessed( - instruction, operand.first, operand.second); - bytes_accessed_from_alternate_mem += operand_bytes_accessed; - } - - for (auto& shape_idx : outputs_in_alternate_mem) { - const float output_bytes_accessed = - op_cost_manager_.OutputBytesAccessed(instruction, shape_idx); - bytes_accessed_from_alternate_mem += output_bytes_accessed; - } - return bytes_accessed_from_alternate_mem; -} - -namespace { -// Returns true on async instructions since we assume they are already -// efficiently scheduled such that they are not in the critical path and appear -// to take no time. -bool ExcludeInstructionFromElapsed( - const HloInstruction& instruction, - const absl::flat_hash_set& execution_threads) { - return !MemorySpaceAssignmentUtils::IsInstructionOnConfiguredExecThread( - instruction, execution_threads) || - instruction.opcode() == HloOpcode::kAllGatherStart || - instruction.opcode() == HloOpcode::kAllGatherDone || - instruction.opcode() == HloOpcode::kAllReduceStart || - instruction.opcode() == HloOpcode::kAllReduceDone || - instruction.opcode() == HloOpcode::kAsyncStart || - instruction.opcode() == HloOpcode::kAsyncDone || - instruction.opcode() == HloOpcode::kCollectivePermuteStart || - instruction.opcode() == HloOpcode::kCollectivePermuteDone || - instruction.opcode() == HloOpcode::kCopyStart || - instruction.opcode() == HloOpcode::kCopyDone; -} -} // namespace - -float CostAnalysis::GetInstructionElapsedDueToCompute( - const HloInstruction& instruction) const { - if (ExcludeInstructionFromElapsed(instruction, options_.execution_threads)) { - return 0.0f; - } - return op_cost_manager_.ComputeSeconds(instruction); -} - -float CostAnalysis::GetInstructionElapsedDueToMemory( - const HloInstruction& instruction, - absl::Span> operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const { - if (ExcludeInstructionFromElapsed(instruction, options_.execution_threads)) { - return 0.0f; - } - float total_bytes_accessed = op_cost_manager_.TotalBytesAccessed(instruction); - float bytes_accessed_read_from_alternate_mem = - GetBytesAccessedFromAlternateMemory(instruction, - operands_in_alternate_mem, {}); - float bytes_accessed_write_in_alternate_mem = - GetBytesAccessedFromAlternateMemory(instruction, {}, - outputs_in_alternate_mem); - - float elapsed_due_to_alternate_mem = - bytes_accessed_read_from_alternate_mem / - options_.alternate_mem_read_bandwidth_bytes_per_second + - bytes_accessed_write_in_alternate_mem / - options_.alternate_mem_write_bandwidth_bytes_per_second; - - float bytes_accessed_from_alternate_mem = - bytes_accessed_read_from_alternate_mem + - bytes_accessed_write_in_alternate_mem; - float elapsed_due_to_default_mem = - (total_bytes_accessed - bytes_accessed_from_alternate_mem) / - DefaultMemBandwidthBytesPerSecond(); - return elapsed_due_to_alternate_mem + elapsed_due_to_default_mem; -} - -float CostAnalysis::GetInstructionElapsedDueToMemory( - const HloInstruction& instruction, - IsInAlternateMemoryFun is_in_alternate_mem) const { - if (ExcludeInstructionFromElapsed(instruction, options_.execution_threads)) { - return 0.0f; - } - float total_bytes_accessed = op_cost_manager_.TotalBytesAccessed(instruction); - float bytes_accessed_from_alternate_mem = 0.0; - for (int operand_num = 0; operand_num < instruction.operand_count(); - ++operand_num) { - ShapeUtil::ForEachSubshape( - instruction.operand(operand_num)->shape(), - [&](const Shape& subshape, const ShapeIndex& index) { - if (!subshape.IsArray()) { - return; - } - if (is_in_alternate_mem(operand_num, index, subshape)) { - bytes_accessed_from_alternate_mem += - op_cost_manager_.OperandBytesAccessed(instruction, operand_num, - index); - } - }); - } - - float bytes_accessed_read_from_alternate_mem = - bytes_accessed_from_alternate_mem; - ShapeUtil::ForEachSubshape(instruction.shape(), [&](const Shape& subshape, - const ShapeIndex& index) { - if (!subshape.IsArray()) { - return; - } - if (is_in_alternate_mem(/*operand_num=*/std::nullopt, index, subshape)) { - bytes_accessed_from_alternate_mem += - op_cost_manager_.OutputBytesAccessed(instruction, index); - } - }); - float bytes_accessed_write_in_alternate_mem = - bytes_accessed_from_alternate_mem - - bytes_accessed_read_from_alternate_mem; - - float elapsed_due_to_alternate_mem = - bytes_accessed_read_from_alternate_mem / - options_.alternate_mem_read_bandwidth_bytes_per_second + - bytes_accessed_write_in_alternate_mem / - options_.alternate_mem_write_bandwidth_bytes_per_second; - float elapsed_due_to_default_mem = - (total_bytes_accessed - bytes_accessed_from_alternate_mem) / - DefaultMemBandwidthBytesPerSecond(); - return elapsed_due_to_alternate_mem + elapsed_due_to_default_mem; -} - -float CostAnalysis::GetInstructionElapsed( - const HloInstruction& instruction) const { - if (ExcludeInstructionFromElapsed(instruction, options_.execution_threads)) { - return 0.0f; - } - float overhead = GetDefaultMemoryAccessOverhead(instruction); - return std::max(GetInstructionElapsedDueToCompute(instruction), - GetInstructionElapsedDueToMemory(instruction) + overhead); -} - -float CostAnalysis::GetInstructionElapsedInAlternateMemory( - const HloInstruction& instruction, - absl::Span> operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const { - if (ExcludeInstructionFromElapsed(instruction, options_.execution_threads)) { - return 0.0f; - } - float overhead = GetDefaultMemoryAccessOverhead( - instruction, operands_in_alternate_mem, outputs_in_alternate_mem); - return std::max( - GetInstructionElapsedDueToCompute(instruction), - GetInstructionElapsedDueToMemory(instruction, operands_in_alternate_mem, - outputs_in_alternate_mem) + - overhead); -} - -float CostAnalysis::GetInstructionElapsedInAlternateMemory( - const HloInstruction& instruction, - IsInAlternateMemoryFun is_in_alternate_mem) const { - if (ExcludeInstructionFromElapsed(instruction, options_.execution_threads)) { - return 0.0f; - } - return std::max( - GetInstructionElapsedDueToCompute(instruction), - GetInstructionElapsedDueToMemory(instruction, is_in_alternate_mem)); -} - -float CostAnalysis::GetAsyncCopyElapsed(int64_t size_in_bytes) const { - return static_cast(size_in_bytes) / - DefaultMemBandwidthBytesPerSecond(/*use_scaling_factor=*/true); -} - -int64_t CostAnalysis::GetScheduleEndTime() const { - return hlo_live_range_->schedule_end_time(); -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h b/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h deleted file mode 100644 index 875f2c575c4d75..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/cost_analysis.h +++ /dev/null @@ -1,268 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_COST_ANALYSIS_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_COST_ANALYSIS_H_ - -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/functional/function_ref.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/call_graph.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -// Options to be passed to the CostAnalysis. -struct CostAnalysisOptions { - // This variable is used by the cost analysis in estimating how many times - // each while loop will execute. Nested loops will be assumed to have - // executed pow(while_execution_count, nesting_level) times. - uint64_t xla_tpu_memory_space_assignment_while_execution_count = 5ULL; - - // This variable is used to scale the alternate memory benefit factor for - // large buffers. The default scaling function is sqrt. - std::string - xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers = - "SQRT"; - - // The window size used to calculate the pipeline overhead when HLO accesses - // the default memory, in MiB. - float pipeline_overhead_window_size_mib = 0; - - double alternate_mem_read_bandwidth_bytes_per_second = 0.0f; - double alternate_mem_write_bandwidth_bytes_per_second = 0.0f; - - double default_mem_bandwidth_bytes_per_second = 0.0f; - - // Scales effective bandwidth for async copies. Valid range is (0, 1]. - float async_copy_bandwidth_scaling_factor = 1.0; - - // Used to get the layout size of a shape in bytes. - std::function shape_size_bytes_fn = - [](const Shape& shape) { return ShapeUtil::ByteSizeOf(shape); }; - - // The execution threads that cost analysis operates on. This is used to - // determine which flattened instructions should be included/excluded from - // cost analysis. - absl::flat_hash_set execution_threads = { - HloInstruction::kMainExecutionThread}; -}; - -// A wrapper class around BaseCosts with additional knowledge about the -// bandwidths of different memory spaces. -class CostAnalysis { - public: - // An optional Cache object may be provided to some of the methods below to - // speed up the lookup. - struct Cache { - // TODO(hanruobing): This map assumes the nested while loops have the same - // hard-coded trip count. We plan to replace it with a more accurate - // estimation provided by 'while_nest_trip_count'. - absl::flat_hash_map while_nest_multiplier; - absl::flat_hash_map computation_trip_count; - absl::flat_hash_map memory_boundedness; - }; - - // Function type that can be used to indicate which input/output values are in - // the alternate memory. - using IsInAlternateMemoryFun = absl::FunctionRef /*operand_num*/, const ShapeIndex& /*index*/, - const Shape& /*shape*/)>; - - virtual ~CostAnalysis() = default; - - static absl::StatusOr> Create( - OpCostManager& op_cost_manager, const CostAnalysisOptions& options, - const AliasInfo* alias_info, const HloModule& module, - HloAliasAnalysis* alias_analysis); - - int64_t GetShapeSizeBytes(const Shape& shape) const; - - float OperandBytesAccessed(const HloInstruction& instruction, - int64_t operand_num, - const ShapeIndex& shape_index) const; - - float OutputBytesAccessed(const HloInstruction& instruction, - const ShapeIndex& shape_index) const; - - double DefaultMemBandwidthBytesPerSecond( - bool use_scaling_factor = false) const; - - // Returns a heuristic value that captures how much putting this tensor to the - // alternate memory would help if the op is memory bound, or otherwise how far - // off is the op to memory boundedness. The larger this number, the higher - // priority it will be placed in the alternate memory. - float GetAlternateMemoryBenefit(const HloInstruction& instruction, - float elapsed_time_due_to_alternate_mem, - Cache* cache = nullptr) const; - // Like above, return the benefit of putting the output tensor in the - // alternate memory. - float GetAlternateMemoryBenefit(const HloPosition& position, - Cache* cache = nullptr) const; - // Like above, return the benefit of putting the input tensor in the alternate - // memory. - float GetAlternateMemoryBenefit(const HloUse& use, - Cache* cache = nullptr) const; - - // Returns a heuristic value of memory boundedness for the given - // BufferInterval. The larger this number, the higher priority it will be - // placed in the alternate memory. - float GetMemoryBoundedness( - const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval, - Cache* cache = nullptr) const; - - // If enabled in CostAnalysisOptions::pipeline_overhead_window_size_mib, - // returns the overhead of accessing the default memory, in seconds. The - // source of the overhead is the software pipelining ovehead. The lowering of - // the operations typically use tiling to copy one window at a time from - // default memory, and perform compute: - // - // Pipeline overhead: <-> - // +----+----+----+----+ - // Copy from default mem: | | | | | - // +----+----+----+----+ - // \ \ \ \ - // \ \ \ \ - // V V V V - // +--+ +--+ +--+ +--+ - // Compute: | | | | | | | | - // +--+ +--+ +--+ +--+ - float GetDefaultMemoryAccessOverhead( - const HloInstruction& instruction, - absl::Span> - operands_in_alternate_mem = {}, - absl::Span outputs_in_alternate_mem = {}) const; - - // Returns the amount of time the default memory bandwidth is idle, while - // executing this instruction, in seconds. This value can be multiplied with - // the default memory bandwidth to get the amount of bytes that are available - // to be copied to/from default memory during the execution of this - // instruction. - float GetDefaultMemoryBandwidthIdleTime( - const HloInstruction& instruction, - absl::Span> - operands_in_alternate_mem = {}, - absl::Span outputs_in_alternate_mem = {}) const; - - // Returns the bytes accessed from alternate memory. - float GetBytesAccessedFromAlternateMemory( - const HloInstruction& instruction, - absl::Span> - operands_in_alternate_mem = {}, - absl::Span outputs_in_alternate_mem = {}) const; - - // Returns the elapsed time in seconds due to compute only. - float GetInstructionElapsedDueToCompute( - const HloInstruction& instruction) const; - - // Returns the elapsed time in seconds due to memory only. If - // operands_in_alternate_mem or outputs_in_alternate_mem is provided, it will - // assume that the corresponding operands or output will be in the alternate - // memory space. This is useful for calculating the benefit of placing the - // buffer in alternate memory. - float GetInstructionElapsedDueToMemory( - const HloInstruction& instruction, - absl::Span> - operands_in_alternate_mem = {}, - absl::Span outputs_in_alternate_mem = {}) const; - - // Like above, only the inputs/outputs indicated by is_in_alternate_mem are in - // the alternate memory. - float GetInstructionElapsedDueToMemory( - const HloInstruction& instruction, - IsInAlternateMemoryFun is_in_alternate_mem) const; - - // Returns the estimated elapsed duration of the instruction in seconds. It - // assumes all operands and outputs of the instruction are in the default - // memory. - virtual float GetInstructionElapsed(const HloInstruction& instruction) const; - - // Returns the estimated elapsed duration of the instruction in seconds. It - // assumes all operands and outputs of the instruction are in the default - // memory, except for the operands and outputs specified to be in the - // alternate memory. - virtual float GetInstructionElapsedInAlternateMemory( - const HloInstruction& instruction, - absl::Span> - operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const; - - // Like above, only the inputs/outputs indicated by is_in_alternate_mem are in - // the alternate memory. - float GetInstructionElapsedInAlternateMemory( - const HloInstruction& instruction, - IsInAlternateMemoryFun is_in_alternate_mem) const; - - // Returns the elapsed time it would take to asynchronously copy the shape - // from default to alternate memory space (or vice versa). - virtual float GetAsyncCopyElapsed(int64_t size_in_bytes) const; - - int64_t GetScheduleEndTime() const; - - // Returns the number of nested computation levels this instruction resides - // in. If while_only is true, it returns the while loop nest level and 0 - // means the instruction is not in a while loop. - int CalculateComputationNestLevel(const HloInstruction* instruction, - bool while_only) const; - - // Returns the number of times the instruction will be executed. - // For instructions in nested loops, this is the product of the number of - // trip counts of outer loops. - float CalculateNestTripCount(const HloInstruction* instruction, - Cache* cache = nullptr) const; - - float GetWhileNestMultiplier(int while_nest_level) const; - - const HloLiveRange& hlo_live_range() const { return *hlo_live_range_; } - - protected: - CostAnalysis(OpCostManager& op_cost_manager, - const CostAnalysisOptions& options, - HloAliasAnalysis* alias_analysis, - std::unique_ptr hlo_live_range, - std::unique_ptr call_graph); - - private: - // A manager responsible for return basic cost metrics. - OpCostManager& op_cost_manager_; - const CostAnalysisOptions options_; - HloAliasAnalysis* alias_analysis_; - std::unique_ptr hlo_live_range_; - std::unique_ptr call_graph_; -}; - -} // namespace memory_space_assignment -} // namespace xla -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_COST_ANALYSIS_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc b/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc deleted file mode 100644 index 603633e6ef1bf6..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc +++ /dev/null @@ -1,297 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/cost_analysis.h" - -#include -#include - -#include -#include -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/status/status_matchers.h" -#include "absl/strings/string_view.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "tsl/platform/errors.h" - -namespace xla { -namespace { - -using memory_space_assignment::CostAnalysis; -using memory_space_assignment::CostAnalysisOptions; - -class MemorySpaceAssignmentCostAnalysisTest - : public HloHardwareIndependentTestBase { - protected: - absl::Status Initialize(const HloModule* module, - float pipeline_overhead_window_size_mib = 0.0) { - HloCostAnalysis::Options options; - options_.alternate_mem_read_bandwidth_bytes_per_second = 128; - options_.alternate_mem_write_bandwidth_bytes_per_second = 128; - options_.default_mem_bandwidth_bytes_per_second = 32; - options_.pipeline_overhead_window_size_mib = - pipeline_overhead_window_size_mib; - options.set_flops_per_second(8); - options.set_bytes_per_second(32); - options.set_transcendentals_per_second(16); - options.set_flops_min_latency_second(1); - hlo_cost_analysis_ = std::make_unique(options); - hlo_cost_analysis_wrapper_ = - std::make_unique(*hlo_cost_analysis_); - op_cost_manager_ = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper_), - /*enable_cache=*/false)); - ABSL_ASSIGN_OR_RETURN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module, &alias_info_)); - ABSL_ASSIGN_OR_RETURN( - cost_analysis_, - CostAnalysis::Create(*op_cost_manager_, options_, &alias_info_, *module, - alias_analysis.get())); - return absl::OkStatus(); - } - - CostAnalysisOptions options_; - std::unique_ptr hlo_cost_analysis_; - std::unique_ptr hlo_cost_analysis_wrapper_; - std::unique_ptr op_cost_manager_; - std::unique_ptr cost_analysis_; - AliasInfo alias_info_; -}; - -TEST_F(MemorySpaceAssignmentCostAnalysisTest, NoPipelineOverhead) { - absl::string_view hlo_string = R"( - HloModule module, is_scheduled=true - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - param1 = f32[2,4] parameter(1) - ROOT add = f32[2,4] add(param0, param1) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK(Initialize(module.get())); - - const HloInstruction* add = module->entry_computation()->root_instruction(); - const float expected_compute_elapsed = std::max( - /*num_flops=*/8.0f / /*flops_per_second=*/8.0f, - hlo_cost_analysis_->min_latency_seconds(HloCostAnalysis::kFlopsKey)); - LOG(INFO) << "Expected compute elapsed = " << expected_compute_elapsed; - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToCompute(*add), - expected_compute_elapsed); - float expected_memory_elapsed = - /*bytes_accessed=*/(3 * 4 * 8) / /*bytes_per_second=*/32.0; - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory(*add), - expected_memory_elapsed); - - // This HLO is memory-bound. - EXPECT_EQ(cost_analysis_->GetInstructionElapsed(*add), - expected_memory_elapsed); - EXPECT_EQ( - cost_analysis_->GetInstructionElapsedInAlternateMemory(*add, {}, {}), - expected_memory_elapsed); - - // Put operand 0 in alternate memory. Still memory bound. - expected_memory_elapsed = - (/*bytes_accessed=*/(2 * 4 * 8) / /*bytes_per_second=*/32.0) + - (/*bytes_accessed=*/(4 * 8) / /*bytes_per_second=*/128.0); - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory(*add, {{0, {}}}), - expected_memory_elapsed); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedInAlternateMemory( - *add, {{0, {}}}, {}), - expected_memory_elapsed); - - // Put operand 0 and output in alternate memory. Still memory bound. - expected_memory_elapsed = - (/*bytes_accessed=*/(4 * 8) / /*bytes_per_second=*/32.0) + - (/*bytes_accessed=*/(2 * 4 * 8) / /*bytes_per_second=*/128.0); - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - EXPECT_EQ( - cost_analysis_->GetInstructionElapsedDueToMemory(*add, {{0, {}}}, {{}}), - expected_memory_elapsed); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedInAlternateMemory( - *add, {{0, {}}}, {{}}), - expected_memory_elapsed); - - // Put everything in alternate memory. We're now compute bound. - expected_memory_elapsed = - /*bytes_accessed=*/(3 * 4 * 8) / /*bytes_per_second=*/128.0; - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory( - *add, {{0, {}}, {1, {}}}, {{}}), - expected_memory_elapsed); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedInAlternateMemory( - *add, {{0, {}}, {1, {}}}, {{}}), - expected_compute_elapsed); -} - -TEST_F(MemorySpaceAssignmentCostAnalysisTest, PipelineOverhead) { - absl::string_view hlo_string = R"( - HloModule module, is_scheduled=true - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - param1 = f32[2,4] parameter(1) - ROOT add = f32[2,4] add(param0, param1) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - // Set the window size 64B. - ASSERT_OK( - Initialize(module.get(), - /*pipeline_overhead_window_size_mib=*/(64.0 / 1024 / 1024))); - - const HloInstruction* add = module->entry_computation()->root_instruction(); - const float expected_compute_elapsed = std::max( - /*num_flops=*/8.0f / /*flops_per_second=*/8.0f, - hlo_cost_analysis_->min_latency_seconds(HloCostAnalysis::kFlopsKey)); - LOG(INFO) << "Expected compute elapsed = " << expected_compute_elapsed; - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToCompute(*add), - expected_compute_elapsed); - float expected_memory_elapsed = - /*bytes_accessed=*/(3 * 4 * 8) / /*bytes_per_second=*/32.0; - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory(*add), - expected_memory_elapsed); - - float expected_overhead = expected_compute_elapsed * 2 / 3; - LOG(INFO) << "Expected overhead = " << expected_overhead; - EXPECT_EQ(cost_analysis_->GetDefaultMemoryAccessOverhead(*add), - expected_overhead); - // This HLO is memory-bound. - EXPECT_EQ(cost_analysis_->GetInstructionElapsed(*add), - expected_memory_elapsed + expected_overhead); - EXPECT_EQ( - cost_analysis_->GetInstructionElapsedInAlternateMemory(*add, {}, {}), - expected_memory_elapsed + expected_overhead); - - // Put operand 0 in alternate memory. Still memory bound. - expected_memory_elapsed = - (/*bytes_accessed=*/(2 * 4 * 8) / /*bytes_per_second=*/32.0) + - (/*bytes_accessed=*/(4 * 8) / /*bytes_per_second=*/128.0); - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - EXPECT_EQ(cost_analysis_->GetDefaultMemoryAccessOverhead(*add, {{0, {}}}), - expected_overhead); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory(*add, {{0, {}}}), - expected_memory_elapsed); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedInAlternateMemory( - *add, {{0, {}}}, {}), - expected_memory_elapsed + expected_overhead); - - // Put operand 0 and output in alternate memory. Still memory bound. - expected_memory_elapsed = - (/*bytes_accessed=*/(4 * 8) / /*bytes_per_second=*/32.0) + - (/*bytes_accessed=*/(2 * 4 * 8) / /*bytes_per_second=*/128.0); - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - expected_overhead = expected_compute_elapsed / 3; - LOG(INFO) << "Expected overhead = " << expected_overhead; - EXPECT_EQ( - cost_analysis_->GetDefaultMemoryAccessOverhead(*add, {{0, {}}}, {{}}), - expected_overhead); - EXPECT_EQ( - cost_analysis_->GetInstructionElapsedDueToMemory(*add, {{0, {}}}, {{}}), - expected_memory_elapsed); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedInAlternateMemory( - *add, {{0, {}}}, {{}}), - expected_memory_elapsed + expected_overhead); - - // Put everything in alternate memory. We're now compute bound. - expected_memory_elapsed = - /*bytes_accessed=*/(3 * 4 * 8) / /*bytes_per_second=*/128.0; - LOG(INFO) << "Expected memory elapsed = " << expected_memory_elapsed; - expected_overhead = 0; - LOG(INFO) << "Expected overhead = " << expected_overhead; - EXPECT_EQ(cost_analysis_->GetDefaultMemoryAccessOverhead( - *add, {{0, {}}, {1, {}}}, {{}}), - expected_overhead); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory( - *add, {{0, {}}, {1, {}}}, {{}}), - expected_memory_elapsed); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedInAlternateMemory( - *add, {{0, {}}, {1, {}}}, {{}}), - expected_compute_elapsed); -} - -TEST_F(MemorySpaceAssignmentCostAnalysisTest, LatencyBoundCompute) { - absl::string_view hlo_string = R"( - HloModule module, is_scheduled=true - - ENTRY Entry { - param0 = f32[2,2] parameter(0) - param1 = f32[2,2] parameter(1) - ROOT add = f32[2,2] add(param0, param1) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK(Initialize(module.get())); - - const HloInstruction* add = module->entry_computation()->root_instruction(); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToCompute(*add), 1.0f); -} - -TEST_F(MemorySpaceAssignmentCostAnalysisTest, ExcludeNonMainThreadComputation) { - absl::string_view hlo_string = R"( - HloModule module, is_scheduled=true - - async_computation { - param0.1 = f32[2,4] parameter(0) - param1.1 = f32[2,4] parameter(1) - ROOT add.1 = f32[2,4] add(param0.1, param1.1) - }, execution_thread="sparse" - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - param1 = f32[2,4] parameter(1) - ROOT add = f32[2,4] add(param0, param1) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK(Initialize(module.get())); - - const HloInstruction* main_add = - module->entry_computation()->root_instruction(); - EXPECT_GT(cost_analysis_->GetInstructionElapsedDueToCompute(*main_add), 0.0f); - - const HloComputation* async_comp = - module->GetComputationWithName("async_computation"); - ASSERT_NE(async_comp, nullptr); - const HloInstruction* async_add = async_comp->root_instruction(); - EXPECT_FALSE(async_add->parent()->IsMainThread()); - EXPECT_EQ(cost_analysis_->GetInstructionElapsed(*async_add), 0.0f); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToCompute(*async_add), - 0.0f); - EXPECT_EQ(cost_analysis_->GetInstructionElapsedDueToMemory( - *async_add, {{0, {}}, {1, {}}}, {{}}), - 0.0f); -} - -} // namespace -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/live_range_util.cc b/third_party/xla/xla/service/memory_space_assignment/live_range_util.cc deleted file mode 100644 index 22005a0798be36..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/live_range_util.cc +++ /dev/null @@ -1,511 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/live_range_util.h" - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/strings/str_cat.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/map_util.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" - -namespace xla { - -namespace { - -const HloValue* FindAssociatedValue(HloPosition pos, const HloBuffer& buffer) { - for (const HloValue* value : buffer.values()) { - for (const HloPosition& p : value->positions()) { - if (p == pos) { - return value; - } - } - } - LOG(FATAL) << "No associated value found for position " << pos.ToString() - << " and buffer " << buffer.ToString(); - return nullptr; -} - -bool IsTrivialOrForwardedPosition(HloPosition position) { - return position.instruction->opcode() == HloOpcode::kGetTupleElement || - position.instruction->opcode() == HloOpcode::kTuple || - position.instruction->opcode() == HloOpcode::kBitcast || - (position.instruction->opcode() == HloOpcode::kAsyncStart && - !position.index.empty() && position.index.front() == 0); -} - -bool IsNonTrivialNonForwardedPosition(HloPosition position) { - return !IsTrivialOrForwardedPosition(position); -} - -HloPosition GetNonTrivialNonForwardedSourcePosition(HloPosition position) { - while (IsTrivialOrForwardedPosition(position)) { - if (position.instruction->opcode() == HloOpcode::kGetTupleElement) { - int64_t tuple_index = position.instruction->tuple_index(); - position.instruction = position.instruction->mutable_operand(0); - position.index.push_front(tuple_index); - } else if (position.instruction->opcode() == HloOpcode::kTuple) { - if (position.index.empty()) { - return position; - } - int64_t tuple_index = position.index.front(); - position.index.pop_front(); - position.instruction = position.instruction->mutable_operand(tuple_index); - } else if (position.instruction->opcode() == HloOpcode::kBitcast) { - position.instruction = position.instruction->mutable_operand(0); - } else if (position.instruction->opcode() == HloOpcode::kAsyncStart && - !position.index.empty() && position.index.front() == 0) { - position.index.pop_front(); - if (position.index.empty()) { - position.instruction = position.instruction->mutable_operand(0); - } else { - position.instruction = - position.instruction->mutable_operand(position.index.front()); - position.index.pop_front(); - } - } - } - return position; -} - -std::string InstructionNameAndComputationName(const HloInstruction* inst) { - return absl::StrCat("(", inst->name(), - " computation: ", inst->parent()->name(), ")"); -} - -bool IsWhileLoopConditionParameter(const HloInstruction* inst) { - if (inst->opcode() != HloOpcode::kParameter || - inst->parent()->caller_instructions().empty()) { - return false; - } - for (const HloInstruction* caller : inst->parent()->caller_instructions()) { - if (caller->opcode() == HloOpcode::kWhile && - caller->while_condition() == inst->parent()) { - return true; - } - } - return false; -} - -} // namespace - -LiveRangeCalculator::LiveRangeCalculator( - const HloBuffer& buffer, - const absl::flat_hash_map& inst_to_index) - : buffer_(buffer), inst_to_index_(inst_to_index) { - // NOLINTNEXTLINE - for (const auto& [inst, index] : inst_to_index) { - index_to_inst_[index] = inst; - auto [it, inserted] = - earliest_inst_time_in_comp_.try_emplace(inst->parent(), index); - if (!inserted) { - it->second = std::min(it->second, index); - } - } - - for (const HloValue* value : buffer_.values()) { - for (const HloUse& use : value->GetUses()) { - if (use.instruction->opcode() == HloOpcode::kAsyncDone) { - // This needs to be updated when async updates are added. - HloInstruction* async_start = use.instruction->async_chain_start(); - int64_t start_time = GetEarliestCalledComputationStartTime(async_start); - async_use_live_ranges_.push_back( - {start_time, inst_to_index.at(use.instruction)}); - } - uses_by_value_and_comp_[value][use.instruction->parent()].push_back(&use); - HloInstruction* operand_instruction = - use.instruction->mutable_operand(use.operand_number); - position_to_uses_[GetNonTrivialNonForwardedSourcePosition(HloPosition{ - operand_instruction, use.operand_index})] - .push_back(&use); - } - } - - async_use_live_ranges_ = SortAndMergeRanges(std::move(async_use_live_ranges_), - /*merge_overlapping_ranges=*/true) - .ranges; -} - -int64_t LiveRangeCalculator::GetInstructionTime( - const HloInstruction* inst) const { - return FindOrDefault(inst_to_index_, inst, -1); -} - -int64_t LiveRangeCalculator::GetEarliestInstructionTime( - const HloComputation* comp) const { - auto it = earliest_inst_time_in_comp_.find(comp); - if (it != earliest_inst_time_in_comp_.end()) { - return it->second; - } - // Fallback to lazy evaluation if the computation wasn't in inst_to_index - int64_t earliest = -1; - for (const HloInstruction* inst : comp->instructions()) { - int64_t time = GetInstructionTime(inst); - if (time != -1) { - if (earliest == -1 || time < earliest) { - earliest = time; - } - } - } - return earliest; -} - -std::vector LiveRangeCalculator::FindLastUsesInComputation( - const HloComputation* comp, const HloValue* value) const { - auto it = uses_by_value_and_comp_.find(value); - if (it == uses_by_value_and_comp_.end()) { - return {}; - } - auto it2 = it->second.find(comp); - if (it2 == it->second.end()) { - return {}; - } - - int64_t max_time = -1; - for (const HloUse* use : it2->second) { - int64_t time = GetInstructionTime(use->instruction); - max_time = std::max(max_time, time); - } - if (max_time == -1) { - return {}; - } - - std::vector last_uses; - for (const HloUse* use : it2->second) { - if (GetInstructionTime(use->instruction) == max_time) { - last_uses.push_back(use); - } - } - return last_uses; -} - -int64_t LiveRangeCalculator::GetEarliestCalledComputationStartTime( - const HloInstruction* inst) const { - int64_t earliest_called_computation_start_time = - std::numeric_limits::max(); - for (const HloComputation* branch_comp : inst->called_computations()) { - earliest_called_computation_start_time = - std::min(earliest_called_computation_start_time, - GetEarliestInstructionTime(branch_comp)); - } - return earliest_called_computation_start_time; -} - -std::vector LiveRangeCalculator::DetermineLiveRangesGivenLastUses( - const std::vector& last_uses, const HloValue* value, - int64_t start_time) { - CHECK(!last_uses.empty()); - HloInstruction* last_inst = last_uses[0]->instruction; - for (const HloUse* last_use : last_uses) { - CHECK_EQ(last_use->instruction, last_inst); - } - HloUse use = *last_uses[0]; - HloInstruction* operand_instruction = - use.instruction->mutable_operand(use.operand_number); - HloPosition current_non_trivial_position = - GetNonTrivialNonForwardedSourcePosition( - HloPosition{operand_instruction, use.operand_index}); - int64_t last_inst_time = GetInstructionTime(last_inst); - std::vector result; - if (last_inst->opcode() != HloOpcode::kWhile && - last_inst->opcode() != HloOpcode::kConditional) { - LiveRange live_range = {start_time, last_inst_time}; - SetLiveRangeForNonTrivialNonForwardedPosition(current_non_trivial_position, - live_range, result); - return result; - } - - if (last_inst->opcode() == HloOpcode::kWhile) { - // While loop output corresponding to the operand index. - HloPosition output_position{use.instruction, use.operand_index}; - // An HloValue that is an operand to a while loop, at one or more operand - // indices, is a pass-through, i.e., it is a read only operand inside the - // while loop. If the last uses at the while loop is a pass-through, the - // live range end would be until the while instruction. - bool pass_through = true; - // If no hlo value in the buffer has a defining position at the - // corresponding while loop output, then the while loop is a pass-through. - for (const HloValue* val : buffer_.values()) { - if (val->defining_instruction() == last_inst) { - pass_through = false; - break; - } - } - if (pass_through) { - LiveRange live_range = {start_time, last_inst_time}; - SetLiveRangeForNonTrivialNonForwardedPosition( - current_non_trivial_position, live_range, result); - return result; - } - } - - // For while loop uses that are not pass-through, the corresponding positions - // inside the condition, body computation and the while loop output are - // considered distinct HloValues. Their live ranges will be computed - // independently. Any special handing is not required. - int64_t earliest_schedule_flattened_computation_start_time = - GetEarliestCalledComputationStartTime(last_inst); - LiveRange live_range = {start_time, - earliest_schedule_flattened_computation_start_time}; - SetLiveRangeForNonTrivialNonForwardedPosition(current_non_trivial_position, - live_range, result); - - if (last_inst->opcode() == HloOpcode::kConditional) { - for (const HloUse* last_use : last_uses) { - std::vector ranges = - DetermineLiveRangesForConditionalUse(last_use, value); - result.insert(result.end(), ranges.begin(), ranges.end()); - } - } - - return SortAndMergeRanges(std::move(result)).ranges; -} - -std::vector LiveRangeCalculator::DetermineLiveRangeInComputation( - HloComputation* comp, const HloValue* value) { - int64_t start_time = GetEarliestInstructionTime(comp); - - // The HloValue should be live from the position inside the computation till - // its last use. There can be multiple uses for the last use instruction. When - // the last use instruction is a special instruction, it can result in - // disjoint live ranges. For example, if the last instruction is a conditional - // instruction. - std::vector internal_last_uses = - FindLastUsesInComputation(comp, value); - if (internal_last_uses.empty()) { - return {{start_time, start_time}}; - } - - return DetermineLiveRangesGivenLastUses(internal_last_uses, value, - start_time); -} - -std::vector -LiveRangeCalculator::DetermineLiveRangesForConditionalUse( - const HloUse* conditional_use, const HloValue* value) { - const HloInstruction* inst = conditional_use->instruction; - CHECK(inst->opcode() == HloOpcode::kConditional); - if (conditional_use->operand_number == 0) { - // Operand 0 of a conditional is the predicate value, it does not have a - // corresponding schedule-flattened computation. - return {}; - } - return DetermineLiveRangeInComputation( - inst->branch_computation(conditional_use->operand_number - 1), value); -} - -LiveRangeResult LiveRangeCalculator::SortAndMergeRanges( - std::vector all_ranges, bool merge_overlapping_ranges) const { - LiveRangeResult result; - if (all_ranges.empty()) { - return result; - } - std::sort(all_ranges.begin(), all_ranges.end(), - [](const LiveRange& a, const LiveRange& b) { - return a.start_time < b.start_time; - }); - std::vector merged; - for (const auto& r : all_ranges) { - if (merged.empty()) { - merged.push_back(r); - continue; - } - if (merged.back().end_time == r.start_time) { - merged.back().end_time = std::max(merged.back().end_time, r.end_time); - } else if (merged.back().end_time < r.start_time) { - merged.push_back(r); - } else { - result.overlap_gt_1_detected = true; - VLOG(1) << "Overlapping live ranges detected: [" - << merged.back().start_time << ", " << merged.back().end_time - << "] and [" << r.start_time << ", " << r.end_time << "]\n" - << "Buffer: " << buffer_.ToString(); - if (merge_overlapping_ranges) { - merged.back().end_time = std::max(merged.back().end_time, r.end_time); - } else { - merged.push_back(r); - } - } - } - result.ranges = merged; - return result; -} - -std::vector LiveRangeCalculator::CalculateHloValueLiveRange( - const HloValue* value) { - VLOG(3) << "Calculating live range for value: " << value->ToString(); - int64_t start_time = GetInstructionTime(value->defining_instruction()); - // - Loop condition parameter positions are considered new HloValues (as - // opposed to forwarded HloPositions). - // - All the HloValues might not have a use in the condition computation, - // but they should be live throughout the computation. - if (IsWhileLoopConditionParameter(value->defining_instruction())) { - int64_t end_time = GetInstructionTime( - value->defining_instruction()->parent()->root_instruction()); - LiveRange live_range = {start_time, end_time}; - std::vector result; - SetLiveRangeForNonTrivialNonForwardedPosition(value->defining_position(), - live_range, result); - return result; - } - - std::vector uses; - uses.reserve(value->GetUses().size()); - for (const HloUse& use : value->GetUses()) { - uses.push_back(&use); - } - if (uses.empty()) { - if (value->defining_instruction()->opcode() == HloOpcode::kFusion) { - // A fusion output HloValue with no uses is a dummy output, which does not - // require any memory allocation. - return {}; - } - std::vector result; - LiveRange live_range = {start_time, start_time}; - SetLiveRangeForNonTrivialNonForwardedPosition(value->defining_position(), - live_range, result); - return result; - } - - std::sort(uses.begin(), uses.end(), [this](const HloUse* a, const HloUse* b) { - return GetInstructionTime(a->instruction) < - GetInstructionTime(b->instruction); - }); - - // The HloValue should be live from definition to the last use instruction. - // There can be multiple uses for the last use instruction. When the last use - // instruction has called computations, it can result in disjoint live ranges. - // For example, if the last instruction is a conditional instruction. - const HloInstruction* last_inst = uses.back()->instruction; - std::vector last_uses; - for (auto it = uses.rbegin(); - it != uses.rend() && (*it)->instruction == last_inst; ++it) { - last_uses.push_back(*it); - } - return DetermineLiveRangesGivenLastUses(last_uses, value, start_time); -} - -LiveRangeResult LiveRangeCalculator::CalculateBufferLiveRange() { - std::vector all_ranges; - all_ranges.reserve(buffer_.values().size()); - for (const HloValue* value : buffer_.values()) { - std::vector val_ranges = CalculateHloValueLiveRange(value); - VLOG(3) << "Live ranges for value " << value->ToString() << ": "; - for (const auto& range : val_ranges) { - VLOG(3) << " " << DetailedLiveRangeDebugString(range); - } - all_ranges.insert(all_ranges.end(), val_ranges.begin(), val_ranges.end()); - } - return SortAndMergeRanges(std::move(all_ranges)); -} - -void LiveRangeCalculator::SetLiveRangeForNonTrivialNonForwardedPosition( - HloPosition pos, LiveRange range, std::vector& live_ranges) { - CHECK(IsNonTrivialNonForwardedPosition(pos)) - << "pos: " << pos.ToString() << " [" << range.start_time << ", " - << range.end_time << "]\nbuffer:" << buffer_.ToString(); - std::optional async_range = - FindContainingAsyncUseLiveRange(range.start_time); - if (async_range.has_value()) { - if (range.end_time <= async_range.value().end_time) { - VLOG(3) << "Non-trivial position " << pos.ToString() << " live range " - << DetailedLiveRangeDebugString(range) - << " is completely within async use live range " - << DetailedLiveRangeDebugString(async_range.value()) - << " and hence not considered.\n"; - return; - } - VLOG(3) << "Non-trivial position " << pos.ToString() << " live range " - << DetailedLiveRangeDebugString(range) - << " is partially within async use live range " - << DetailedLiveRangeDebugString(async_range.value()) - << ". Updating live range and position."; - range.start_time = std::max(range.start_time, async_range.value().end_time); - const HloValue* associated_value = FindAssociatedValue(pos, buffer_); - - for (const HloPosition& p : associated_value->positions()) { - if (GetInstructionTime(p.instruction) == async_range.value().end_time && - IsNonTrivialNonForwardedPosition(p)) { - pos = p; - break; - } - } - VLOG(3) << "Updated non-trivial position " << pos.ToString() - << " live range " << DetailedLiveRangeDebugString(range) << "\n"; - } - auto live_range_it = position_to_live_range_.find(pos); - if (live_range_it != position_to_live_range_.end()) { - CHECK_EQ(live_range_it->second.start_time, range.start_time); - CHECK_EQ(live_range_it->second.end_time, range.end_time); - return; - } - - VLOG(3) << "Setting live range for non-trivial position " << pos.ToString() - << " to " << DetailedLiveRangeDebugString(range) << "\n"; - position_to_live_range_[pos] = range; - live_ranges.push_back(range); -} - -std::optional LiveRangeCalculator::FindContainingAsyncUseLiveRange( - int64_t start_time) const { - auto it = std::lower_bound(async_use_live_ranges_.begin(), - async_use_live_ranges_.end(), start_time, - [](const LiveRange& element, int64_t value) { - return element.end_time < value; - }); - if (it != async_use_live_ranges_.end()) { - const auto& async_range = *it; - if (async_range.start_time <= start_time && - start_time <= async_range.end_time) { - return async_range; - } - } - return std::nullopt; -} - -std::optional -LiveRangeCalculator::GetLiveRangeForNonTrivialNonForwardedPosition( - HloPosition pos) const { - auto it = position_to_live_range_.find(pos); - if (it != position_to_live_range_.end()) { - return it->second; - } - return std::nullopt; -} - -std::string LiveRangeCalculator::DetailedLiveRangeDebugString( - const LiveRange& range) const { - const HloInstruction* start_inst = index_to_inst_.at(range.start_time); - const HloInstruction* end_inst = index_to_inst_.at(range.end_time); - return absl::StrCat( - range.ToString(), - " start_inst: ", InstructionNameAndComputationName(start_inst), - " end_inst: ", InstructionNameAndComputationName(end_inst)); -} - -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/live_range_util.h b/third_party/xla/xla/service/memory_space_assignment/live_range_util.h deleted file mode 100644 index 775d0f2c650890..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/live_range_util.h +++ /dev/null @@ -1,241 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_LIVE_RANGE_UTIL_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_LIVE_RANGE_UTIL_H_ - -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/strings/str_cat.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" - -namespace xla { - -// Represents a closed interval [start_time, end_time] of instruction schedule -// indices during which a value or buffer is live. -struct LiveRange { - int64_t start_time; - int64_t end_time; - - std::string ToString() const { - return absl::StrCat("[", start_time, ", ", end_time, "]"); - } -}; - -struct LiveRangeResult { - std::vector ranges; - // Two live ranges [a, b] and [c, d] are considered overlapping only if b > c, - // if b == c, they are merged, because one HloPosition can start at the same - // instuction as another HloPosition ends (e.g., dynamic-update-slice). If - // overlap_gt_1_detected is true, it only means that overlapping live ranges - // were detected, it does not mean that overlaps were merged/removed. - bool overlap_gt_1_detected = false; -}; - -// Calculates the live ranges of HloValues and HloBuffers based on a flattened -// instruction schedule. It handles schedule-flattened computations (like -// branches in conditionals and bodies of while loops) to accurately determine -// the span of liveness across control flow boundaries. Schedule-flattened -// computations impose constraints regarding the continuity of buffer liveness -// across boundaries. Additionally, async ops, loops, and conditionals can -// introduce additional complexity. -// -// Below we use the term "last uses" of an HloValue, which raises the question, -// how can we have more than 1 last use. There are 2 ways: -// - An HloValue may be used multiple times by the same instruction, e.g., -// add(a, a). -// - When dealing with conditionals, each branch has its own last use of an -// HloValue. -class LiveRangeCalculator { - public: - // Constructs a calculator for the given buffer, using the provided - // instruction-to-index mapping to determine time. - LiveRangeCalculator( - const HloBuffer& buffer, - const absl::flat_hash_map& inst_to_index); - - // Returns a merged and sorted list of live ranges for the entire buffer - // (union of live ranges of all its values). Also returns whether overlapping - // live ranges were detected during sorting and merging. Note: Overlapping - // live ranges are not merged/removed. - LiveRangeResult CalculateBufferLiveRange(); - - // Trivial positions are not eligible for real memory allocations, they are - // GTE, tuple, bitcast, or forwarded operand positions in an async - // instruction. If the live range of a non-trivial position of `buffer_` is - // not a sub-interval of the live range of any other non-trivial position of - // the same buffer, it is considered an "outer" position, and it requires a - // real memory allocation. This function returns the live range of the given - // non-trivial position, if it is an "outer" position and the live range is - // already calculated. Otherwise, it returns nullopt. - // NOTE: CalculateBufferLiveRange() (or CalculateHloValueLiveRange()) must be - // executed prior to querying individual position live ranges. - std::optional GetLiveRangeForNonTrivialNonForwardedPosition( - HloPosition pos) const; - - // Returns a debug string representation of the live range, including - // the names of the start and end instructions. - std::string DetailedLiveRangeDebugString(const LiveRange& range) const; - - // Returns the earliest flattened-schedule index of any instruction within the - // given computation. - int64_t GetEarliestInstructionTime(const HloComputation* comp) const; - - const absl::flat_hash_map& - non_trivial_non_forwarded_position_to_live_range() const { - return position_to_live_range_; - } - - const absl::flat_hash_map>& - non_trivial_non_forwarded_position_to_uses() const { - return position_to_uses_; - } - - // Sorts the given live ranges by start time and merges overlapping ranges. - LiveRangeResult SortAndMergeRanges( - std::vector all_ranges, - bool merge_overlapping_ranges = false) const; - - // Returns a merged and sorted list of live ranges for the given HLO value, - // starting from its definition to its last use(s). - std::vector CalculateHloValueLiveRange(const HloValue* value); - - // Returns the schedule index (time) of the given instruction. - // Returns -1 if the instruction is not in the schedule. - int64_t GetInstructionTime(const HloInstruction* inst) const; - - private: - // Returns the last (with highest schedule index) uses of the given value in - // the given computation. The last uses returned by this function might not be - // true use sites (where a value is read/used), e.g., when the last use is a - // conditional. This method is not recursive. - std::vector FindLastUsesInComputation( - const HloComputation* comp, const HloValue* value) const; - - // For instructions that call computations, that are inlined in the flattened - // schedule (e.g., While, Conditional), returns the earliest start time, among - // all schedule-flattened computations. Note: The instruction itself is not - // considered; it is scheduled after the schedule-flattened computations. - int64_t GetEarliestCalledComputationStartTime( - const HloInstruction* inst) const; - - // Determines the live ranges for a value given its last uses. - // If the last use is a control flow instruction (While/Conditional), - // it recursively determines liveness within the schedule-flattened - // computations. Note: kWhile and kConditional instructions are handled - // differently from other instructions with schedule-flattened computations. - // REQUIRES: `last_uses` is not empty, and all uses in `last_uses` must belong - // to the same instruction. - std::vector DetermineLiveRangesGivenLastUses( - const std::vector& last_uses, const HloValue* value, - int64_t start_time); - - // Determines the live range of a value within a specific schedule-flattened - // computation (e.g., a branch of a conditional or body of a loop). - std::vector DetermineLiveRangeInComputation(HloComputation* comp, - const HloValue* value); - - // Determines the live ranges of a value for a use in a Conditional - // instruction. - std::vector DetermineLiveRangesForConditionalUse( - const HloUse* conditional_use, const HloValue* value); - - // Associates a calculated live range with a non-trivial position. - void SetLiveRangeForNonTrivialNonForwardedPosition( - HloPosition pos, LiveRange range, std::vector& live_ranges); - - // Returns the async use live range that contains the given start time, if - // one exists. - std::optional FindContainingAsyncUseLiveRange( - int64_t start_time) const; - - // The buffer for which we are calculating live ranges. - const HloBuffer& buffer_; - - // A map from instruction to its index in the flattened schedule. - const absl::flat_hash_map& inst_to_index_; - - // A reverse mapping from instruction index in the flattened schedule to - // instruction. - absl::flat_hash_map index_to_inst_; - - // A map to store the earliest instruction time in each computation. - absl::flat_hash_map - earliest_inst_time_in_comp_; - - // A nested hash map to find all the uses of a value in a computation. - // Format: uses_by_value_and_comp_[value][computation] = vector_of_uses. - absl::flat_hash_map< - const HloValue*, - absl::flat_hash_map>> - uses_by_value_and_comp_; - - // A map from non-trivial HloPositions to all the uses of that position. - // The HloUse objects are owned by HloValue (which in turn are owned by - // HloAliasAnalysis). This map stores non-owning raw pointers; the calculator - // must not outlive the HloAliasAnalysis. - absl::flat_hash_map> - position_to_uses_; - - // Contains live ranges of non-trivial positions of the buffer that are not - // sub-intervals of the live range of any other non-trivial position of the - // same buffer. From an MSA perspective, these "outer" positions require - // independent memory allocations, whereas the nested "inner" positions - // can share/mirror the allocation of their containing position. - absl::flat_hash_map position_to_live_range_; - - // Live ranges for which buffer_ is required to have contiguous allocation - // because of async uses. If an async operand has an aliased use inside the - // async called computation, it will result in overlapping live ranges because - // the new HloValue inside the async called computation is live at the same - // time as the async operand. This is safe to merge and does not qualify as a - // real overlap. This data structure is used to identify such cases. - // - // Example: - // - // async_comp (p0: f32[2,2]) -> f32[2,2] { - // p2 = f32[2,2] parameter(1) - // ROOT custom_call0 = f32[2,2] custom-call(p2) out_to_op_alias={{0}:(0)} - // } - // - // ENTRY entry (p0: f32[2,2], p1: f32[2,2]) -> f32[2,2] { - // p0 = f32[2,2] parameter(0) - // p1 = f32[2,2] parameter(1) - // as_start = {f32[2,2], f32[2,2], s32[]} async-start(p0) calls=async_comp - // negate0 = f32[2,2] negate(p1) - // async_done = f32[2,2] async-done(as_start) - // ROOT add0 = f32[2,2] add(negate0, async_done) - // } - // - // In the example above, the HloPosition as_start {0} is a forwarded - // position of p0 and requires a contiguous allocation and thus has a live - // range from as_start to async_done. The async computation - // has an aliased use of p0, which defines a new HloValue for the same buffer, - // that has a live range from custom_call0 to add0. This results in an - // overlap, which is legal and safe to merge. - std::vector async_use_live_ranges_; -}; - -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_LIVE_RANGE_UTIL_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/live_range_util_test.cc b/third_party/xla/xla/service/memory_space_assignment/live_range_util_test.cc deleted file mode 100644 index 44b3e0b1d8d5c1..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/live_range_util_test.cc +++ /dev/null @@ -1,795 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/live_range_util.h" - -#include -#include -#include -#include - -#include -#include -#include "absl/container/flat_hash_map.h" -#include "absl/log/log.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_schedule.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/hlo_buffer.h" -#include "xla/shape_util.h" - -namespace xla { -namespace { - -class LiveRangeCalculatorTest : public HloHardwareIndependentTestBase { - protected: - absl::StatusOr> - GetInstToIndex(const HloModule* module, - const HloAliasAnalysis& alias_analysis) { - ABSL_ASSIGN_OR_RETURN(auto live_range, - HloLiveRange::Run(module->schedule(), alias_analysis, - module->entry_computation())); - return live_range->instruction_schedule(); - } - - const HloBuffer* GetBufferForInstruction( - const HloAliasAnalysis& alias_analysis, const HloInstruction* inst, - const ShapeIndex& index = {}) { - return &alias_analysis.GetUniqueBufferAt(inst, index); - } -}; - -// Tests live range calculation for a simple linear HLO sequence without -// control flow. -TEST_F(LiveRangeCalculatorTest, SimpleLinear) { - const char* hlo_text = R"( - HloModule Simple, is_scheduled=true - ENTRY %entry (p0: f32[2,3], p1: f32[2,3]) -> f32[2,3] { - %p0 = f32[2,3]{1,0} parameter(0) - %p1 = f32[2,3]{1,0} parameter(1) - %add = f32[2,3]{1,0} add(%p0, %p1) - ROOT %neg = f32[2,3]{1,0} negate(%add) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* add_inst = FindInstruction(module.get(), "add"); - ASSERT_NE(add_inst, nullptr); - const HloBuffer* add_buffer = - GetBufferForInstruction(*alias_analysis, add_inst); - ASSERT_NE(add_buffer, nullptr); - - LiveRangeCalculator calculator(*add_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p0 - // 1: p1 - // 2: add - // 3: neg - - ASSERT_EQ(ranges.size(), 1); - EXPECT_EQ(ranges[0].start_time, 2); - EXPECT_EQ(ranges[0].end_time, 3); -} - -// Tests Case 1: Conditional input that has no subsequent use after the -// conditional. The live range should propagate into the branches and result in -// disjoint live ranges for each branch. -TEST_F(LiveRangeCalculatorTest, ConditionalInputNoSubsequentUse) { - const char* hlo_text = R"( - HloModule Cond, is_scheduled=true - true_comp (p3: (f32[2,3])) -> (f32[2,3]) { - p3 = (f32[2,3]) parameter(0) - gte1 = f32[2,3] get-tuple-element(p3), index=0 - neg0 = f32[2,3] negate(gte1) - ROOT tuple2 = (f32[2,3]) tuple(neg0) - } - false_comp (p4: (f32[2,3])) -> (f32[2,3]) { - p4 = (f32[2,3]) parameter(0) - gte2 = f32[2,3] get-tuple-element(p4), index=0 - neg1 = f32[2,3] negate(gte2) - ROOT tuple3 = (f32[2,3]) tuple(neg1) - } - ENTRY entry (p0: pred[], p1: f32[2,3]) -> f32[2,3] { - p0 = pred[] parameter(0) - p1 = f32[2,3] parameter(1) - tanh0 = f32[2,3] tanh(p1) - tuple0 = (f32[2,3]) tuple(tanh0) - conditional = (f32[2,3]) conditional(p0, tuple0, tuple0), true_computation=true_comp, false_computation=false_comp - gte3 = f32[2,3] get-tuple-element(conditional), index=0 - ROOT neg2 = f32[2,3] negate(gte3) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p0 - // 1: p1 - // 2: tanh0 - // 3: tuple0 - // -- true_comp -- - // 4: p3 - // 5: gte1 - // 6: neg0 - // 7: tuple2 - // -- false_comp -- - // 8: p4 - // 9: gte2 - // 10: neg1 - // 11: tuple3 - // -- entry -- - // 12: conditional - // 13: gte3 - // 14: neg2 - - // tanh0 buffer should be live: - // - in entry: from tanh0 (2) to tuple0 (3) -> [2, 3] - // - in true_comp: from p3 (4) to neg0 (6) -> [4, 6] - // - in false_comp: from p4 (8) to neg1 (10) -> [8, 10] - // Merged ranges should be: [2, 6] and [8, 10] - // Wait, tanh0 value itself is live [2, 4] and [8, 8] (due to conditional at - // 12, earliest called comp is 4, so [2, 4]. And false_comp starts at 8, so - // [8, 8]). p3 index 0 is live [4, 6]. p4 index 0 is live [8, 10]. Union: [2, - // 4] U [8, 8] U [4, 6] U [8, 10] = [2, 6] and [8, 10]. - - ASSERT_EQ(ranges.size(), 2); - EXPECT_EQ(ranges[0].start_time, 2); - EXPECT_EQ(ranges[0].end_time, 6); - EXPECT_EQ(ranges[1].start_time, 8); - EXPECT_EQ(ranges[1].end_time, 10); -} - -// Tests Case 2: Conditional input that has a subsequent use after the -// conditional. The live range should cover the entire region from definition to -// the subsequent use. -TEST_F(LiveRangeCalculatorTest, ConditionalInputWithSubsequentUse) { - const char* hlo_text = R"( - HloModule Cond, is_scheduled=true - true_comp (p3: (f32[2,3])) -> (f32[2,3]) { - p3 = (f32[2,3]) parameter(0) - gte1 = f32[2,3] get-tuple-element(p3), index=0 - neg0 = f32[2,3] negate(gte1) - ROOT tuple2 = (f32[2,3]) tuple(neg0) - } - false_comp (p4: (f32[2,3])) -> (f32[2,3]) { - p4 = (f32[2,3]) parameter(0) - gte2 = f32[2,3] get-tuple-element(p4), index=0 - neg1 = f32[2,3] negate(gte2) - ROOT tuple3 = (f32[2,3]) tuple(neg1) - } - ENTRY entry (p0: pred[], p1: f32[2,3]) -> f32[2,3] { - p0 = pred[] parameter(0) - p1 = f32[2,3] parameter(1) - tanh0 = f32[2,3] tanh(p1) - tuple0 = (f32[2,3]) tuple(tanh0) - conditional = (f32[2,3]) conditional(p0, tuple0, tuple0), true_computation=true_comp, false_computation=false_comp - gte3 = f32[2,3] get-tuple-element(conditional), index=0 - neg2 = f32[2,3] negate(gte3) - add0 = f32[2,3] add(tanh0, neg2) - ROOT neg3 = f32[2,3] negate(add0) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p0 - // 1: p1 - // 2: tanh0 - // 3: tuple0 - // -- true_comp -- - // 4: p3 - // 5: gte1 - // 6: neg0 - // 7: tuple2 - // -- false_comp -- - // 8: p4 - // 9: gte2 - // 10: neg1 - // 11: tuple3 - // -- entry -- - // 12: conditional - // 13: gte3 - // 14: neg2 - // 15: add0 - // 16: neg3 - - // tanh0 buffer should be live: - // - from tanh0 (2) to add0 (15) because of subsequent use at add0. - // - inside true_comp: [4, 6] (p3 to neg0) - // - inside false_comp: [8, 10] (p4 to neg1) - // Union of [2, 15], [4, 6], [8, 10] is [2, 15]. - - ASSERT_EQ(ranges.size(), 1); - EXPECT_EQ(ranges[0].start_time, 2); - EXPECT_EQ(ranges[0].end_time, 15); -} - -// Tests Case 3: Conditional output where branch outputs and conditional output -// are mapped to the same buffer, resulting in separate, non-overlapping live -// ranges. -TEST_F(LiveRangeCalculatorTest, ConditionalOutput) { - const char* hlo_text = R"( - HloModule Cond, is_scheduled=true - true_comp (p3: (f32[2,3])) -> (f32[2,3]) { - p3 = (f32[2,3]) parameter(0) - gte1 = f32[2,3] get-tuple-element(p3), index=0 - neg0 = f32[2,3] negate(gte1) - ROOT tuple2 = (f32[2,3]) tuple(neg0) - } - false_comp (p4: (f32[2,3])) -> (f32[2,3]) { - p4 = (f32[2,3]) parameter(0) - gte2 = f32[2,3] get-tuple-element(p4), index=0 - neg1 = f32[2,3] negate(gte2) - ROOT tuple3 = (f32[2,3]) tuple(neg1) - } - ENTRY entry (p0: pred[], p1: f32[2,3]) -> f32[2,3] { - p0 = pred[] parameter(0) - p1 = f32[2,3] parameter(1) - tanh0 = f32[2,3] tanh(p1) - tuple0 = (f32[2,3]) tuple(tanh0) - conditional = (f32[2,3]) conditional(p0, tuple0, tuple0), true_computation=true_comp, false_computation=false_comp - gte3 = f32[2,3] get-tuple-element(conditional), index=0 - ROOT neg2 = f32[2,3] negate(gte3) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* cond_inst = - FindInstruction(module.get(), "conditional"); - ASSERT_NE(cond_inst, nullptr); - const HloBuffer* cond_output_buffer = - GetBufferForInstruction(*alias_analysis, cond_inst, {0}); - ASSERT_NE(cond_output_buffer, nullptr); - - LiveRangeCalculator calculator(*cond_output_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p0 - // 1: p1 - // 2: tanh0 - // 3: tuple0 - // -- true_comp -- - // 4: p3 - // 5: gte1 - // 6: neg0 - // 7: tuple2 - // -- false_comp -- - // 8: p4 - // 9: gte2 - // 10: neg1 - // 11: tuple3 - // -- entry -- - // 12: conditional - // 13: gte3 - // 14: neg2 - - // The conditional output buffer (element 0) contains: - // - neg0 (def 6) in true_comp, used as ROOT. Live range: [6, 7] (tuple2 is at - // 7) - // - neg1 (def 10) in false_comp, used as ROOT. Live range: [10, 11] (tuple3 - // is at 11) - // - conditional output (def 12) in entry, used at neg2 (14). Live range: [12, - // 14] Since they don't merge (7 < 8, 11 < 12), we expect 3 ranges: [6, 7], - // [10, 11], [12, 14]. - - ASSERT_EQ(ranges.size(), 3); - EXPECT_EQ(ranges[0].start_time, 6); - EXPECT_EQ(ranges[0].end_time, 7); - EXPECT_EQ(ranges[1].start_time, 10); - EXPECT_EQ(ranges[1].end_time, 11); - EXPECT_EQ(ranges[2].start_time, 12); - EXPECT_EQ(ranges[2].end_time, 14); -} - -// Tests Case 4: Conditional input that aliases the output, merging them into a -// single continuous live range. -TEST_F(LiveRangeCalculatorTest, ConditionalInputAliasOutput) { - const char* hlo_text = R"( - HloModule Cond, is_scheduled=true - true_comp (p3: (f32[2,3])) -> (f32[2,3]) { - p3 = (f32[2,3]) parameter(0) - gte1 = f32[2,3] get-tuple-element(p3), index=0 - update = f32[1,1] constant({{1}}) - index = s32[] constant(0) - dus = f32[2,3] dynamic-update-slice(gte1, update, index, index) - ROOT tuple2 = (f32[2,3]) tuple(dus) - } - false_comp (p4: (f32[2,3])) -> (f32[2,3]) { - p4 = (f32[2,3]) parameter(0) - gte2 = f32[2,3] get-tuple-element(p4), index=0 - update = f32[1,1] constant({{1}}) - index = s32[] constant(0) - dus = f32[2,3] dynamic-update-slice(gte2, update, index, index) - ROOT tuple3 = (f32[2,3]) tuple(dus) - } - ENTRY entry (p0: pred[], p1: f32[2,3]) -> f32[2,3] { - p0 = pred[] parameter(0) - p1 = f32[2,3] parameter(1) - tanh0 = f32[2,3] tanh(p1) - tuple0 = (f32[2,3]) tuple(tanh0) - conditional = (f32[2,3]) conditional(p0, tuple0, tuple0), true_computation=true_comp, false_computation=false_comp - gte3 = f32[2,3] get-tuple-element(conditional), index=0 - ROOT neg2 = f32[2,3] negate(gte3) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - const HloInstruction* cond_inst = - FindInstruction(module.get(), "conditional"); - ASSERT_NE(cond_inst, nullptr); - const HloBuffer* cond_output_buffer = - GetBufferForInstruction(*alias_analysis, cond_inst, {0}); - ASSERT_EQ(tanh0_buffer, cond_output_buffer); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p0 - // 1: p1 - // 2: tanh0 - // 3: tuple0 - // -- true_comp -- - // 4: p3 - // 5: gte1 - // 6: update (true) - // 7: index (true) - // 8: dus (true) - // 9: tuple2 - // -- false_comp -- - // 10: p4 - // 11: gte2 - // 12: update (false) - // 13: index (false) - // 14: dus (false) - // 15: tuple3 - // -- entry -- - // 16: conditional - // 17: gte3 - // 18: neg2 - - // - [2, 9] covers tanh0 definition to true_comp end. - // - [10, 15] covers false_comp start to false_comp end. - // - [16, 18] covers conditional output to neg2. - ASSERT_EQ(ranges.size(), 3); - EXPECT_EQ(ranges[0].start_time, 2); - EXPECT_EQ(ranges[0].end_time, 9); - EXPECT_EQ(ranges[1].start_time, 10); - EXPECT_EQ(ranges[1].end_time, 15); - EXPECT_EQ(ranges[2].start_time, 16); - EXPECT_EQ(ranges[2].end_time, 18); -} - -// Tests live range calculation for loop-carried parameters in a while loop, -// resulting in disjoint live ranges for loop iterations. -TEST_F(LiveRangeCalculatorTest, WhileLoop) { - const char* hlo_text = R"( - HloModule While, is_scheduled=true - cond_comp (p0: (f32[2,3], s32[])) -> pred[] { - p0 = (f32[2,3], s32[]) parameter(0) - limit = s32[] constant(10) - iteration = s32[] get-tuple-element(p0), index=1 - ROOT cond = pred[] compare(iteration, limit), direction=LT - } - body_comp (p1: (f32[2,3], s32[])) -> (f32[2,3], s32[]) { - p1 = (f32[2,3], s32[]) parameter(0) - val = f32[2,3] get-tuple-element(p1), index=0 - iteration = s32[] get-tuple-element(p1), index=1 - one = s32[] constant(1) - new_iteration = s32[] add(iteration, one) - new_val = f32[2,3] negate(val) - ROOT tuple2 = (f32[2,3], s32[]) tuple(new_val, new_iteration) - } - ENTRY entry (p2: f32[2,3], p3: s32[]) -> (f32[2,3], s32[]) { - p2 = f32[2,3] parameter(0) - p3 = s32[] parameter(1) - tanh0 = f32[2,3] tanh(p2) - init_tuple = (f32[2,3], s32[]) tuple(tanh0, p3) - ROOT while_loop = (f32[2,3], s32[]) while(init_tuple), condition=cond_comp, body=body_comp - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p2 - // 1: p3 - // 2: tanh0 - // 3: init_tuple - // -- cond_comp -- - // 4: p0 - // 5: limit - // 6: iteration - // 7: cond - // -- body_comp -- - // 8: p1 - // 9: val - // 10: iteration - // 11: one - // 12: new_iteration - // 13: new_val - // 14: tuple2 - // -- entry -- - // 15: while_loop - - // Expected ranges: - // - tanh0 (def 2) to start of loop comps (4) -> [2, 4] - // - p0 index 0 (def 4) -> [4, 4] (no uses) - // - p1 index 0 (def 8) to new_val (13) -> [8, 13] - // - new_val (def 13) to tuple2 (14) -> [13, 14] - // - while_loop output element 0 (def 15) -> [15, 15] - // Merged: [2, 4], [8, 14], [15, 15] - - ASSERT_EQ(ranges.size(), 3); - EXPECT_EQ(ranges[0].start_time, 2); - EXPECT_EQ(ranges[0].end_time, 7); - EXPECT_EQ(ranges[1].start_time, 8); - EXPECT_EQ(ranges[1].end_time, 14); - EXPECT_EQ(ranges[2].start_time, 15); - EXPECT_EQ(ranges[2].end_time, 15); -} - -// Tests live range calculation for a pass-through HLO value in a while loop, -// where the while instruction is the last use of the pass-through value. -TEST_F(LiveRangeCalculatorTest, WhileLoopPassThrough) { - const char* hlo_text = R"( - HloModule While, is_scheduled=true - cond_comp (p0: (f32[2,3], f32[2,3], s32[])) -> pred[] { - p0 = (f32[2,3], f32[2,3], s32[]) parameter(0) - limit = s32[] constant(10) - iteration = s32[] get-tuple-element(p0), index=2 - ROOT cond = pred[] compare(iteration, limit), direction=LT - } - body_comp (p1: (f32[2,3], f32[2,3], s32[])) -> (f32[2,3], f32[2,3], s32[]) { - p1 = (f32[2,3], f32[2,3], s32[]) parameter(0) - val0 = f32[2,3] get-tuple-element(p1), index=0 - val1 = f32[2,3] get-tuple-element(p1), index=1 - iteration = s32[] get-tuple-element(p1), index=2 - new_val1 = f32[2,3] negate(val1) - one = s32[] constant(1) - new_iteration = s32[] add(iteration, one) - ROOT tuple2 = (f32[2,3], f32[2,3], s32[]) tuple(val0, new_val1, new_iteration) - } - ENTRY entry (p2: f32[2,3], p3: f32[2,3], p4: s32[]) -> f32[2,3] { - p2 = f32[2,3] parameter(0) - p3 = f32[2,3] parameter(1) - p4 = s32[] parameter(2) - tanh0 = f32[2,3] tanh(p2) - tanh1 = f32[2,3] tanh(p3) - init_tuple = (f32[2,3], f32[2,3], s32[]) tuple(tanh0, tanh1, p4) - while_loop = (f32[2,3], f32[2,3], s32[]) while(init_tuple), condition=cond_comp, body=body_comp - gte_updated = f32[2,3] get-tuple-element(while_loop), index=1 - ROOT root_neg = f32[2,3] negate(gte_updated) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - const HloInstruction* while_inst = - FindInstruction(module.get(), "while_loop"); - ASSERT_NE(while_inst, nullptr); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p2 - // 1: p3 - // 2: p4 - // 3: tanh0 - // 4: tanh1 - // 5: init_tuple - // -- cond_comp -- - // 6: p0 - // 7: limit - // 8: iteration - // 9: cond - // -- body_comp -- - // 10: p1 - // 11: val0 - // 12: val1 - // 13: iteration - // 14: new_val1 - // 15: one - // 16: new_iteration - // 17: tuple2 - // -- entry -- - // 18: while_loop - // 19: gte_updated - // 20: root_neg - - // Because tanh0 is pass-through in the while loop and while_loop is its last - // use, its live range extends continuously from its definition (tanh0) until - // the while_loop instruction. - ASSERT_EQ(ranges.size(), 1); - EXPECT_EQ(ranges[0].start_time, inst_to_index.at(tanh0_inst)); - EXPECT_EQ(ranges[0].end_time, inst_to_index.at(while_inst)); - EXPECT_EQ(ranges[0].start_time, 3); - EXPECT_EQ(ranges[0].end_time, 18); -} - -// Tests live range calculation for async instructions, ensuring the input -// buffer remains live from async-start until async-done. -TEST_F(LiveRangeCalculatorTest, AsyncInstruction) { - const char* hlo_text = R"( - HloModule Async, is_scheduled=true - async_comp (p0: f32[2,3]) -> f32[2,3] { - p0 = f32[2,3] parameter(0) - ROOT neg0 = f32[2,3] negate(p0) - } - ENTRY entry (p1: f32[2,3]) -> f32[2,3] { - p1 = f32[2,3] parameter(0) - tanh0 = f32[2,3] tanh(p1) - async_start = ((f32[2,3]), f32[2,3], s32[]) async-start(tanh0), calls=async_comp - async_done = f32[2,3] async-done(async_start) - ROOT neg1 = f32[2,3] negate(async_done) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p1 - // 1: tanh0 - // 2: p0 - // 3: neg0 - // 4: async_start - // 5: async_done - // 6: neg1 - - // The input buffer must remain live from tanh0 (1) until async_done (5) - // completes. - ASSERT_EQ(ranges.size(), 1); - EXPECT_EQ(ranges[0].start_time, 1); - EXPECT_EQ(ranges[0].end_time, 5); -} - -TEST_F(LiveRangeCalculatorTest, NestedConditional) { - const char* hlo_text = R"( - HloModule NestedCond, is_scheduled=true - inner_true (p5: (f32[2,3])) -> (f32[2,3]) { - p5 = (f32[2,3]) parameter(0) - gte_in = f32[2,3] get-tuple-element(p5), index=0 - neg_in = f32[2,3] negate(gte_in) - ROOT tuple_in = (f32[2,3]) tuple(neg_in) - } - - inner_false (p6: (f32[2,3])) -> (f32[2,3]) { - p6 = (f32[2,3]) parameter(0) - gte_in2 = f32[2,3] get-tuple-element(p6), index=0 - abs_in = f32[2,3] abs(gte_in2) - ROOT tuple_in2 = (f32[2,3]) tuple(abs_in) - } - - outer_true (p3: (f32[2,3], pred[])) -> (f32[2,3]) { - p3 = (f32[2,3], pred[]) parameter(0) - gte_out = f32[2,3] get-tuple-element(p3), index=0 - pred_out = pred[] get-tuple-element(p3), index=1 - tuple_out = (f32[2,3]) tuple(gte_out) - cond_in = (f32[2,3]) conditional(pred_out, tuple_out, tuple_out), true_computation=inner_true, false_computation=inner_false - gte_out2 = f32[2,3] get-tuple-element(cond_in), index=0 - ROOT tuple_out2 = (f32[2,3]) tuple(gte_out2) - } - - outer_false (p4: (f32[2,3])) -> (f32[2,3]) { - p4 = (f32[2,3]) parameter(0) - gte_out3 = f32[2,3] get-tuple-element(p4), index=0 - neg_out = f32[2,3] negate(gte_out3) - ROOT tuple_out3 = (f32[2,3]) tuple(neg_out) - } - - ENTRY entry (p0: pred[], p1: f32[2,3], p2: pred[]) -> f32[2,3] { - p0 = pred[] parameter(0) - p1 = f32[2,3] parameter(1) - p2 = pred[] parameter(2) - tanh0 = f32[2,3] tanh(p1) - tuple0 = (f32[2,3], pred[]) tuple(tanh0, p2) - tuple1 = (f32[2,3]) tuple(tanh0) - conditional = (f32[2,3]) conditional(p0, tuple0, tuple1), true_computation=outer_true, false_computation=outer_false - gte_entry = f32[2,3] get-tuple-element(conditional), index=0 - ROOT neg_entry = f32[2,3] negate(gte_entry) - } - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_TRUE(module->has_schedule()); - - AliasInfo alias_info; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info)); - ASSERT_OK_AND_ASSIGN(auto inst_to_index, - GetInstToIndex(module.get(), *alias_analysis)); - - const HloInstruction* tanh0_inst = FindInstruction(module.get(), "tanh0"); - ASSERT_NE(tanh0_inst, nullptr); - const HloBuffer* tanh0_buffer = - GetBufferForInstruction(*alias_analysis, tanh0_inst); - ASSERT_NE(tanh0_buffer, nullptr); - - LiveRangeCalculator calculator(*tanh0_buffer, inst_to_index); - auto ranges = calculator.CalculateBufferLiveRange().ranges; - - // Expected flattened schedule: - // 0: p0 - // 1: p1 - // 2: p2 - // 3: tanh0 - // 4: tuple0 - // 5: tuple1 - // -- outer_true -- - // 6: p3 - // 7: gte_out - // 8: pred_out - // 9: tuple_out - // -- inner_true -- - // 10: p5 - // 11: gte_in - // 12: neg_in - // 13: tuple_in - // -- inner_false -- - // 14: p6 - // 15: gte_in2 - // 16: abs_in - // 17: tuple_in2 - // -- outer_true (cont) -- - // 18: cond_in - // 19: gte_out2 - // 20: tuple_out2 - // -- outer_false -- - // 21: p4 - // 22: gte_out3 - // 23: neg_out - // 24: tuple_out3 - // -- entry (cont) -- - // 25: conditional - // 26: gte_entry - // 27: neg_entry - - // tanh0 buffer should be live: - // - in entry: [3, 6] (tanh0 (3) -> conditional (25), earliest called comp is - // outer_true (6)) - // - in outer_true: [6, 10] (p3 (6) -> cond_in (18), earliest called comp is - // inner_true (10)) - // - in inner_true: [10, 12] (p5 (10) -> neg_in (12)) - // - in inner_false: [14, 16] (p6 (14) -> abs_in (16)) - // - in outer_false: [21, 23] (p4 (21) -> neg_out (23)) - // Merged ranges: [3, 12], [14, 16], [21, 23] - ASSERT_EQ(ranges.size(), 3); - EXPECT_EQ(ranges[0].start_time, 3); - EXPECT_EQ(ranges[0].end_time, 12); - EXPECT_EQ(ranges[1].start_time, 14); - EXPECT_EQ(ranges[1].end_time, 16); - EXPECT_EQ(ranges[2].start_time, 21); - EXPECT_EQ(ranges[2].end_time, 23); -} - -} // namespace -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc deleted file mode 100644 index 011d4adce60e6e..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc +++ /dev/null @@ -1,1630 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/memory_bound_loop_optimizer.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/base/nullability.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/check.h" -#include "absl/log/die_if_null.h" -#include "absl/log/log.h" -#include "absl/memory/memory.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_format.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/hlo_alias_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/utils/hlo_live_range.h" -#include "xla/service/buffer_value.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo.pb.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -struct LoopOptimizerChunkInterval { - int64_t begin_idx_in_loop; - int64_t end_idx_in_loop; - EvenOddChunkPair chunks; - - std::string ToString() const { - CHECK(chunks.HasValues()); - return absl::StrFormat( - "begin_idx_in_loop: %d, end_idx_in_loop: %d, even chunk: %s, odd " - "chunk: %s", - begin_idx_in_loop, end_idx_in_loop, chunks.even_chunk->ToString(), - chunks.odd_chunk->ToString()); - } -}; - -std::optional GetInstructionIndex( - const HloInstruction* instruction, - const absl::flat_hash_map& - instructions_to_index) { - auto it = instructions_to_index.find(instruction); - return it == instructions_to_index.end() ? std::nullopt - : std::optional(it->second); -} - -} // namespace - -void LoopOptimizerBestFitHeap::CreateBufferInterval( - const AllocationBlock& allocation_block, - const AllocationBlock* colocated_with) { - buffer_intervals_[&allocation_block] = - BufferInterval({&allocation_block, - allocation_block.size, - allocation_block.inclusive_start_time, - allocation_block.end_time, - {}, - colocated_with == nullptr}); - if (colocated_with) { - buffer_intervals_[colocated_with].colocations.push_back(&allocation_block); - } -} - -std::optional -LoopOptimizerBestFitHeap::MaybeFindChunkCandidate( - const AllocationBlock& allocation_block, int64_t preferred_offset) { - Chunk chunk_candidate = FindChunkCandidate( - buffer_intervals_[&allocation_block], preferred_offset); - if (chunk_candidate.chunk_end() <= size_limit_per_heap_) { - return chunk_candidate; - } - return std::nullopt; -} - -std::optional -LoopOptimizerBestFitHeap::FindAndCommitChunkCandidate( - const AllocationBlock& allocation_block, int64_t preferred_offset) { - std::optional chunk = - MaybeFindChunkCandidate(allocation_block, preferred_offset); - if (chunk.has_value()) { - CommitChunkAndInterval(buffer_intervals_[&allocation_block], chunk.value()); - } - return chunk; -} - -void LoopOptimizerBestFitHeap::RemoveChunk(int64_t start_time, int64_t end_time, - Chunk chunk) { - CHECK(interval_tree_.Remove(start_time, end_time, chunk)); -} - -void LoopOptimizerBestFitHeap::RemoveEvenChunks( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, - std::optional& chunk) { - RemoveChunk(begin_idx_in_loop, end_idx_in_loop, chunk.value()); - RemoveChunk(begin_idx_in_loop + 2 * loop_size_, - end_idx_in_loop + 2 * loop_size_, chunk.value()); -} - -void LoopOptimizerBestFitHeap::RemoveOddChunks( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, - std::optional& chunk) { - RemoveChunk(begin_idx_in_loop + loop_size_, end_idx_in_loop + loop_size_, - chunk.value()); - RemoveChunk(begin_idx_in_loop + 3 * loop_size_, - end_idx_in_loop + 3 * loop_size_, chunk.value()); -} - -void LoopOptimizerBestFitHeap::RemoveEvenOddChunkPair( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, - EvenOddChunkPair& chunks) { - CheckAllocationIntervalValid(begin_idx_in_loop, end_idx_in_loop); - ShiftAllocationIntervalIfRequired(begin_idx_in_loop, end_idx_in_loop); - auto& [even_chunk, odd_chunk] = chunks; - RemoveEvenChunks(begin_idx_in_loop, end_idx_in_loop, even_chunk); - RemoveOddChunks(begin_idx_in_loop, end_idx_in_loop, odd_chunk); -} - -const AllocationBlock& LoopOptimizerBestFitHeap::GetAllocationBlock( - int64_t start_time, int64_t end_time, int64_t size) { - allocation_blocks_.push_back( - {start_time, end_time, size, static_cast(-1), - static_cast(-1), - static_cast(allocation_blocks_.size())}); - return allocation_blocks_.back(); -} - -const AllocationBlock& LoopOptimizerBestFitHeap::CreateEvenAllocationBlock( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size) { - const AllocationBlock& first_allocation_block = - GetAllocationBlock(begin_idx_in_loop, end_idx_in_loop, size); - CreateBufferInterval(first_allocation_block); - const AllocationBlock& second_allocation_block = - GetAllocationBlock(begin_idx_in_loop + 2 * loop_size_, - end_idx_in_loop + 2 * loop_size_, size); - CreateBufferInterval(second_allocation_block, &first_allocation_block); - return first_allocation_block; -} - -const AllocationBlock& LoopOptimizerBestFitHeap::CreateOddAllocationBlock( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size) { - const AllocationBlock& first_allocation_block = GetAllocationBlock( - begin_idx_in_loop + loop_size_, end_idx_in_loop + loop_size_, size); - CreateBufferInterval(first_allocation_block); - const AllocationBlock& second_allocation_block = - GetAllocationBlock(begin_idx_in_loop + 3 * loop_size_, - end_idx_in_loop + 3 * loop_size_, size); - CreateBufferInterval(second_allocation_block, &first_allocation_block); - return first_allocation_block; -} - -void LoopOptimizerBestFitHeap::CheckAllocationIntervalValid( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop) const { - CHECK_LE(begin_idx_in_loop, end_idx_in_loop); - CHECK_LE(-1 * loop_size_, begin_idx_in_loop); - CHECK_LT(begin_idx_in_loop, loop_size_); - CHECK_LE(0, end_idx_in_loop); - CHECK_LT(end_idx_in_loop, 2 * loop_size_); - CHECK_LE(end_idx_in_loop - begin_idx_in_loop + 1, 2 * loop_size_); -} - -void LoopOptimizerBestFitHeap::ShiftAllocationIntervalIfRequired( - int64_t& begin_idx_in_loop, int64_t& end_idx_in_loop) const { - if (begin_idx_in_loop < 0) { - begin_idx_in_loop += loop_size_; - end_idx_in_loop += loop_size_; - } -} - -EvenOddChunkPair LoopOptimizerBestFitHeap::FindEvenAndOddAllocationBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - std::pair preferred_offsets) { - CheckAllocationIntervalValid(begin_idx_in_loop, end_idx_in_loop); - ShiftAllocationIntervalIfRequired(begin_idx_in_loop, end_idx_in_loop); - auto [even_offset, odd_offset] = preferred_offsets; - const AllocationBlock& even_allocation = - CreateEvenAllocationBlock(begin_idx_in_loop, end_idx_in_loop, size); - const AllocationBlock& odd_allocation = - CreateOddAllocationBlock(begin_idx_in_loop, end_idx_in_loop, size); - // We need to commit the even chunk because even and odd chunks might overlap - // in time. - std::optional even_chunk = - FindAndCommitChunkCandidate(even_allocation, even_offset); - if (!even_chunk.has_value()) { - return {std::nullopt, std::nullopt}; - } - std::optional odd_chunk = - MaybeFindChunkCandidate(odd_allocation, odd_offset); - RemoveEvenChunks(begin_idx_in_loop, end_idx_in_loop, even_chunk); - if (odd_chunk.has_value()) { - return {even_chunk, odd_chunk}; - } - return {std::nullopt, std::nullopt}; -} - -EvenOddChunkPair LoopOptimizerBestFitHeap::AllocateEvenAndOddBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - std::pair preferred_offsets) { - CheckAllocationIntervalValid(begin_idx_in_loop, end_idx_in_loop); - ShiftAllocationIntervalIfRequired(begin_idx_in_loop, end_idx_in_loop); - auto [even_offset, odd_offset] = preferred_offsets; - const AllocationBlock& even_allocation = - CreateEvenAllocationBlock(begin_idx_in_loop, end_idx_in_loop, size); - const AllocationBlock& odd_allocation = - CreateOddAllocationBlock(begin_idx_in_loop, end_idx_in_loop, size); - // We need to commit the even chunk because even and odd chunks might overlap - // in time. - std::optional even_chunk = - FindAndCommitChunkCandidate(even_allocation, even_offset); - if (!even_chunk.has_value()) { - return {std::nullopt, std::nullopt}; - } - std::optional odd_chunk = - FindAndCommitChunkCandidate(odd_allocation, odd_offset); - if (odd_chunk.has_value()) { - return {even_chunk, odd_chunk}; - } - // Remove even chunk if odd chunk was not found. - RemoveEvenChunks(begin_idx_in_loop, end_idx_in_loop, even_chunk); - return {std::nullopt, std::nullopt}; -} - -const AllocationBlock& -LoopOptimizerBestFitHeap::CreateSameEvenAndOddAllocationBlock( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size) { - const AllocationBlock& first_allocation_block = - GetAllocationBlock(begin_idx_in_loop, end_idx_in_loop, size); - CreateBufferInterval(first_allocation_block); - const AllocationBlock& second_allocation_block = - GetAllocationBlock(begin_idx_in_loop + 1 * loop_size_, - end_idx_in_loop + 1 * loop_size_, size); - CreateBufferInterval(second_allocation_block, &first_allocation_block); - const AllocationBlock& third_allocation_block = - GetAllocationBlock(begin_idx_in_loop + 2 * loop_size_, - end_idx_in_loop + 2 * loop_size_, size); - CreateBufferInterval(third_allocation_block, &first_allocation_block); - const AllocationBlock& fourth_allocation_block = - GetAllocationBlock(begin_idx_in_loop + 3 * loop_size_, - end_idx_in_loop + 3 * loop_size_, size); - CreateBufferInterval(fourth_allocation_block, &first_allocation_block); - return first_allocation_block; -} - -EvenOddChunkPair LoopOptimizerBestFitHeap::FindSameEvenAndOddAllocationBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - int64_t preferred_offset) { - CheckAllocationIntervalValid(begin_idx_in_loop, end_idx_in_loop); - ShiftAllocationIntervalIfRequired(begin_idx_in_loop, end_idx_in_loop); - // An allocation that is colocated in even and odd iterations cannot be double - // buffered i.e. it should span less than or equal to one loop iteration). - CHECK_LE(end_idx_in_loop - begin_idx_in_loop + 1, loop_size_); - const AllocationBlock& allocation = CreateSameEvenAndOddAllocationBlock( - begin_idx_in_loop, end_idx_in_loop, size); - std::optional chunk = - MaybeFindChunkCandidate(allocation, preferred_offset); - return {chunk, chunk}; -} - -EvenOddChunkPair LoopOptimizerBestFitHeap::AllocateSameEvenAndOddBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - int64_t preferred_offset) { - CheckAllocationIntervalValid(begin_idx_in_loop, end_idx_in_loop); - ShiftAllocationIntervalIfRequired(begin_idx_in_loop, end_idx_in_loop); - // An allocation that is colocated in even and odd iterations cannot be double - // buffered i.e. it should span less than or equal to one loop iteration). - CHECK_LE(end_idx_in_loop - begin_idx_in_loop + 1, loop_size_); - const AllocationBlock& allocation = CreateSameEvenAndOddAllocationBlock( - begin_idx_in_loop, end_idx_in_loop, size); - std::optional chunk = - FindAndCommitChunkCandidate(allocation, preferred_offset); - return {chunk, chunk}; -} - -std::string LoopOptimizerBestFitHeap::MemoryUsageToAsciiArt( - int64_t begin_iteration, int64_t end_iteration) const { - CHECK_LE(0, begin_iteration); - CHECK_LE(begin_iteration, end_iteration); - return interval_tree_.NodesOverlappingInTimeToAsciiArt( - loop_size_ * begin_iteration, loop_size_ * (end_iteration + 1) - 1, - loop_size_); -} - -std::vector LoopOptimizerBestFitHeap::RemainingMemoryByTime() const { - // Only 2nd and 3rd iterations have the correct (and identical) memory usage. - // 1st and 4th iterations serve only to model the boundary conditions. - std::vector memory_used_by_time = - interval_tree_.MemoryUsedInInterval(loop_size_ * 2, loop_size_ * 3 - 1); - std::vector remaining_memory_by_time(loop_size_); - for (int i = 0; i < loop_size_; ++i) { - remaining_memory_by_time[i] = size_limit_per_heap_ - memory_used_by_time[i]; - } - return remaining_memory_by_time; -} - -int64_t LoopOptimizerBestFitHeap::LastMemoryOffsetOccupied() const { - // 2nd and 3rd iterations will suffice for getting the current alternate - // memory size. - return interval_tree_.HeapSizeInInterval(loop_size_ * 2, loop_size_ * 4 - 1); -} - -/*static*/ absl::StatusOr> -MemoryBoundLoopOptimizer::Create(int loop_start, int loop_end, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis, - const Options& options) { - CHECK(options.cost_analysis != nullptr); - std::unique_ptr optimizer = - absl::WrapUnique(new MemoryBoundLoopOptimizer( - loop_start, loop_end, options.max_size_in_bytes, - options.memory_bound_loop_optimizer_options, hlo_live_range, - alias_analysis, *options.cost_analysis, &options.size_fn, - options.reserved_scoped_memory_fn, options.alignment_in_bytes)); - ABSL_RETURN_IF_ERROR(optimizer->Initialize()); - return std::move(optimizer); -} - -MemoryBoundLoopOptimizer::MemoryBoundLoopOptimizer( - int loop_start, int loop_end, uint64_t alternate_memory_size, - const MemoryBoundLoopOptimizerOptions& options, - const HloLiveRange& hlo_live_range, const HloAliasAnalysis& alias_analysis, - const CostAnalysis& cost_analysis, - const BufferValue::SizeFunction* absl_nonnull size_function, - const ReservedScopedMemoryFunction& reserved_scoped_memory_fn, - int64_t alignment_in_bytes) - : loop_start_(loop_start), - loop_end_(loop_end), - loop_size_(loop_end - loop_start), - alternate_memory_size_(alternate_memory_size), - options_(options), - hlo_live_range_(hlo_live_range), - alias_analysis_(alias_analysis), - cost_analysis_(cost_analysis), - size_function_(ABSL_DIE_IF_NULL(size_function)), - reserved_scoped_memory_fn_(reserved_scoped_memory_fn), - heap_(LoopOptimizerBestFitHeap(alternate_memory_size, - /*loop_size=*/loop_end - loop_start, - alignment_in_bytes)) {} - -absl::Status MemoryBoundLoopOptimizer::Initialize() { - const auto& instruction_sequence = - hlo_live_range_.flattened_instruction_sequence().instructions(); - VLOG(3) << "MemoryBoundLoopOptimizer::Initialize, loop start: " << loop_start_ - << ", loop end: " << loop_end_ << ", loop size: " << loop_size_ - << ", alternate memory size: " << alternate_memory_size_; - const HloComputation* loop_computation = nullptr; - // Initialize the remaining memory array with the size of the alternate - // memory. Also populate instructions_in_loop_ and - // instructions_in_{prev,next}_iterations_ data structures to help find the - // loop values. - int prev_iteration_start = loop_start_ - loop_size_; - int next_iteration_start = loop_start_ + loop_size_; - for (int i = 0; i < loop_size_; ++i) { - const HloInstruction* loop_inst = instruction_sequence[loop_start_ + i]; - instructions_in_loop_[loop_inst] = i; - const HloInstruction* prev_iteration_inst = - instruction_sequence[prev_iteration_start + i]; - instructions_in_prev_iteration_[prev_iteration_inst] = i; - const HloInstruction* next_iteration_inst = - instruction_sequence[next_iteration_start + i]; - instructions_in_next_iteration_[next_iteration_inst] = i; - - VLOG(3) << " inst in loop [" << (i) << "]: " << loop_inst->name(); - if (!loop_computation) { - loop_computation = loop_inst->parent(); - } else { - TF_RET_CHECK(loop_computation == loop_inst->parent()); - } - int64_t reserved_memory = - reserved_scoped_memory_fn_(loop_inst, - /*operands_in_alternate_memory=*/{}, - /*outputs_in_alternate_memory=*/{}); - if (reserved_memory == 0) { - continue; - } - // Chunks for reserved scoped memory should always be found at offset 0. - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - i, i, reserved_memory, /*preferred_offsets=*/{0, 0}); - CHECK(chunks.HasValues()); - CHECK(chunks.even_chunk->size == reserved_memory); - VLOG(3) << "Reserved chunk: " << chunks.even_chunk->ToString() - << " loop index: " << i; - } - - // Create a tree set to keep track of all the values that the loop - // instructions produce and consume. We use a tree set instead of a hash set - // to ensure the iteration order is the same as insertion order. Since we - // traverse the program in instruction order, the buffers would be inserted in - // a deterministic order, so we'll be able to iterate over these buffers in a - // deterministic order. - std::set buffers_to_process; - for (const auto& [instruction, idx] : instructions_in_loop_) { - auto maybe_add_buffer = [&](const HloInstruction* instruction) { - return [this, &buffers_to_process, instruction](const Shape& subshape, - const ShapeIndex& index) { - if (!subshape.IsArray()) { - return; - } - const HloBuffer& buffer = - alias_analysis_.GetUniqueBufferAt(instruction, index); - if (buffers_to_process.find(&buffer) == buffers_to_process.end()) { - buffers_to_process.insert(&buffer); - } - }; - }; - ShapeUtil::ForEachSubshape(instruction->shape(), - maybe_add_buffer(instruction)); - for (const HloInstruction* operand : instruction->operands()) { - ShapeUtil::ForEachSubshape(operand->shape(), maybe_add_buffer(operand)); - } - } - - // Process the buffers and decide if they should be added as LoopValues. - for (const HloBuffer* buffer : buffers_to_process) { - MaybeCreateLoopValue(*buffer, loop_computation); - } - return absl::OkStatus(); -} - -void MemoryBoundLoopOptimizer::MaybeCreateLoopValue( - const HloBuffer& buffer, const HloComputation* loop_computation) { - loop_values_.push_back({}); - LoopValue& loop_value = loop_values_.back(); - float pos_bytes = 0; - float use_bytes = 0; - bool has_footer_consumer = false; - for (const HloValue* value : buffer.values()) { - // For each position and use of the value, populate the respective position - // and use fields for the current, previous, and next iterations along with - // the loop indices. - for (const HloPosition& position : value->positions()) { - if (position.instruction->opcode() == HloOpcode::kGetTupleElement || - position.instruction->opcode() == HloOpcode::kTuple) { - continue; - } - std::optional loop_index = - GetInstructionIndex(position.instruction, instructions_in_loop_); - std::optional prev_iteration_index; - if (loop_index) { - loop_value.loop_positions.push_back({*loop_index, position}); - VLOG(3) << "Pos match: " << position.instruction->name() << " at " - << *loop_index; - } else if ((prev_iteration_index = GetInstructionIndex( - position.instruction, instructions_in_prev_iteration_))) { - loop_value.prev_iteration_positions.push_back( - {*prev_iteration_index, position}); - VLOG(3) << "Pos match (prev iteration): " - << position.instruction->name() << " at " - << *prev_iteration_index; - } else if (loop_value.prev_iteration_positions.empty() && - loop_value.loop_positions.empty() && - position.instruction->parent() == loop_computation && - !loop_value.header_position) { - loop_value.header_position = position; - } - - // Keep track of bytes accessed by this value. - if (loop_index || prev_iteration_index) { - float bytes_accessed = cost_analysis_.OutputBytesAccessed( - *position.instruction, position.index); - pos_bytes += bytes_accessed; - VLOG(3) << " accessed: " << bytes_accessed; - } - } - - for (const HloUse& use : value->GetUses()) { - if (use.instruction->opcode() == HloOpcode::kGetTupleElement) { - continue; - } - std::optional loop_index = - GetInstructionIndex(use.instruction, instructions_in_loop_); - std::optional next_iteration_index; - if (loop_index) { - loop_value.loop_uses.push_back({*loop_index, use}); - VLOG(3) << "Use match: " << use.instruction->name() << " at " - << *loop_index; - } else if ((next_iteration_index = GetInstructionIndex( - use.instruction, instructions_in_next_iteration_))) { - loop_value.next_iteration_uses.push_back({*next_iteration_index, use}); - VLOG(3) << "Use match (next iteration): " << use.instruction->name() - << " at " << *next_iteration_index; - } else if (!loop_value.loop_positions.empty() || - !loop_value.loop_uses.empty()) { - has_footer_consumer = true; - } - - // Keep track of bytes accessed by this value. - if (loop_index || next_iteration_index) { - float bytes_accessed = cost_analysis_.OperandBytesAccessed( - *use.instruction, use.operand_number, use.operand_index); - use_bytes += bytes_accessed; - VLOG(3) << " accessed: " << bytes_accessed; - } - } - } - - // We only add the loop position if it has a position or use in the current - // iteration and its previous iteration positions are empty. The reason why we - // disallow values with previous iteration positions is because there will be - // a different value that corresponds to the same value but one iteration - // later, so we will add that one instead. - if ((!loop_value.loop_positions.empty() || !loop_value.loop_uses.empty()) && - loop_value.prev_iteration_positions.empty()) { - loop_value.size = (*size_function_)(**buffer.values().begin()); - VLOG(3) << "Size: " << loop_value.size; - // Classify the type of allocation. See the comment in LoopValue definition. - loop_value.allocation_type = LoopValue::AllocationType::kUnsupported; - auto position_compare = [](const std::pair& a, - const std::pair& b) { - return a.first < b.first; - }; - auto use_compare = [](const std::pair& a, - const std::pair& b) { - return a.first < b.first; - }; - absl::c_sort(loop_value.loop_positions, position_compare); - absl::c_sort(loop_value.prev_iteration_positions, position_compare); - absl::c_sort(loop_value.loop_uses, use_compare); - absl::c_sort(loop_value.next_iteration_uses, use_compare); - if (!loop_value.loop_positions.empty()) { - if (loop_value.next_iteration_uses.empty() && - !loop_value.loop_uses.empty()) { - loop_value.allocation_type = LoopValue::AllocationType::kTemporary; - } else if (!loop_value.next_iteration_uses.empty()) { - if (loop_value.next_iteration_uses.back().first >= - loop_value.loop_positions.front().first) { - loop_value.allocation_type = - LoopValue::AllocationType::kLoopCarriedDependence; - } else { - loop_value.allocation_type = LoopValue::AllocationType::kTemporary; - } - } - } else if (loop_value.header_position && !loop_value.loop_uses.empty()) { - if (loop_value.loop_uses.size() == - loop_value.next_iteration_uses.size() && - loop_value.loop_uses.front().first == - loop_value.next_iteration_uses.front().first) { - loop_value.allocation_type = LoopValue::AllocationType::kPinned; - } else if (loop_value.next_iteration_uses.empty() || - loop_value.next_iteration_uses.back().first < - loop_value.loop_uses.front().first) { - loop_value.allocation_type = LoopValue::AllocationType::kPrefetch; - } - } - - VLOG(3) << "Allocation type " - << LoopValue::AllocationTypeToString(loop_value.allocation_type); - VLOG(3) << "Pos bytes: " << pos_bytes << " use bytes: " << use_bytes; - - // We calculate the savings of allocating this buffer in the alternate - // memory. - float savings = pos_bytes + use_bytes; - if (loop_value.header_position) { - savings -= loop_value.size; - } - if (!loop_value.loop_positions.empty() && has_footer_consumer) { - savings -= loop_value.size; - } - loop_value.savings = savings; - loop_value.savings_per_byte = savings / loop_value.size; - VLOG(3) << "Savings: " << loop_value.savings; - VLOG(3) << "Savings per byte: " << loop_value.savings_per_byte; - for (const HloValue* value : buffer.values()) { - VLOG(6) << value->ToString(); - } - loop_value.hlo_values = buffer.values(); - } else { - loop_values_.pop_back(); - } -} - -void MemoryBoundLoopOptimizer::Optimize() { - SortLoopValues(); - AllocateLoopValues(); - PostProcess(); -} - -float MemoryBoundLoopOptimizer::CalculateExecutionTime() const { - // First populate the list of prefetches. - std::vector> prefetches; - for (const LoopValue& value : loop_values_) { - if (!value.allocations.empty() && - value.allocations.back()->is_copy_allocation()) { - prefetches.push_back( - {static_cast(value.allocations.back().get()), - cost_analysis_.GetAsyncCopyElapsed(cost_analysis_.GetShapeSizeBytes( - value.hlo_values.front()->shape()))}); - } - } - - // Returns the effective prefetch completion time. The effective time is a - // value that will be larger than loop size for prefetches that start in this - // iteration but complete in the next iteration. - auto get_effective_done_time = - [&](int64_t copy_start_schedule_after, - int64_t copy_done_schedule_before) -> int64_t { - if (copy_start_schedule_after == loop_size_ - 1 && - copy_done_schedule_before == 0) { - return 2 * loop_size_; - } - if (copy_start_schedule_after + 1 >= copy_done_schedule_before) { - return copy_done_schedule_before + loop_size_; - } - return copy_done_schedule_before; - }; - - // Sort the prefetches by first the start time, then the effective done time. - absl::c_sort( - prefetches, [&](const std::pair& a, - const std::pair& b) { - return std::forward_as_tuple( - a.first->copy_start_schedule_after(), - get_effective_done_time( - a.first->copy_start_schedule_after(), - a.first->copy_done_schedule_before())) < - std::forward_as_tuple(b.first->copy_start_schedule_after(), - get_effective_done_time( - b.first->copy_start_schedule_after(), - b.first->copy_done_schedule_before())); - }); - // Populate the required prefetch completions array. For each instruction in - // the loop, this vector holds the index of the latest-issued prefetch that - // needs to be completed before the instruction executes, or nullopt if there - // is no prefetch that needs to finish by this instruction. To represent - // prefetches that started in the previous iteration, we use negative numbers. - std::vector> required_prefetch_completions(loop_size_); - for (int i = 0; i < prefetches.size(); ++i) { - const auto& [prefetch, elapsed] = prefetches[i]; - int required_prefetch_completion = i; - if (prefetch->copy_start_schedule_after() == loop_size_ - 1 && - prefetch->copy_done_schedule_before() == 0) { - required_prefetch_completion -= 2 * prefetches.size(); - } else if (prefetch->copy_start_schedule_after() + 1 >= - prefetch->copy_done_schedule_before()) { - required_prefetch_completion -= prefetches.size(); - } - VLOG(3) << "Prefetch #" << i << " (elapsed " << elapsed - << "): " << prefetch->ToString(); - if (required_prefetch_completions[prefetch->copy_done_schedule_before()]) { - required_prefetch_completions[prefetch->copy_done_schedule_before()] = - std::max( - *required_prefetch_completions[prefetch - ->copy_done_schedule_before()], - required_prefetch_completion); - } else { - required_prefetch_completions[prefetch->copy_done_schedule_before()] = - required_prefetch_completion; - } - VLOG(4) - << "Required completion at " << prefetch->copy_done_schedule_before() - << " = " - << *required_prefetch_completions[prefetch - ->copy_done_schedule_before()]; - } - - // Populate the elapsed times of instructions and bandwidth idle times at each - // point. - float result; - std::vector bandwidth_idle_times; - std::vector instructions_elapsed; - bandwidth_idle_times.reserve(loop_size_); - instructions_elapsed.reserve(loop_size_); - for (int i = 0; i < loop_size_; ++i) { - bandwidth_idle_times.push_back(GetBandwidthIdleTime(i)); - instructions_elapsed.push_back(GetInstructionElapsed(i)); - } - // We simulate the loop for three iterations to measure the steady state. - const int kNumIterations = 3; - // This data structure keeps track of the elapsed time remaining of each - // prefetch. Note that there is a separate entry for each prefetch in each - // iteration simulated. - std::vector prefetch_remaining_elapsed_times(prefetches.size() * - kNumIterations); - int prefetch_start_index = 0; - int prefetch_done_index = 0; - int prefetch_completed_index = 0; - - for (int iteration = 0; iteration < kNumIterations; ++iteration) { - float total_elapsed = 0; - float total_bandwidth_idle_time = 0; - float total_critical_prefetch = 0; - for (int i = 0; i < loop_size_; ++i) { - // If any prefetches are expected to be completed, check if they have any - // remaining elapsed time associated with them, and if so add this to - // critical prefetch time. - std::optional required_prefetch_completion = - required_prefetch_completions[i]; - if (required_prefetch_completion) { - int required_prefetch_done_index = - iteration * static_cast(prefetches.size()) + - *required_prefetch_completion; - VLOG(4) << "Prefetch #" - << ((*required_prefetch_completion + prefetches.size()) % - prefetches.size()) - << " (" << required_prefetch_done_index - << ") is required to be completed at " << i; - for (; prefetch_done_index <= required_prefetch_done_index; - ++prefetch_done_index) { - CHECK_LE(prefetch_done_index, prefetch_start_index); - if (prefetch_done_index == prefetch_completed_index) { - float& prefetch_remaining = - prefetch_remaining_elapsed_times[prefetch_done_index]; - VLOG(4) << "Prefetch #" << (prefetch_done_index % prefetches.size()) - << " (" << prefetch_done_index - << ") did not complete, remaining elapsed = " - << prefetch_remaining; - total_critical_prefetch += prefetch_remaining; - prefetch_remaining = 0; - ++prefetch_completed_index; - } - } - } - - float elapsed = instructions_elapsed[i]; - total_elapsed += elapsed; - float bandwidth_idle_time = bandwidth_idle_times[i]; - // Find the outstanding prefetches during this instruction, and if any of - // them have remaining time, spend some or all of the bandwidth idle time - // to satisfy them. - for (; prefetch_completed_index < prefetch_start_index; - ++prefetch_completed_index) { - float& prefetch_remaining = - prefetch_remaining_elapsed_times[prefetch_completed_index]; - if (bandwidth_idle_time < prefetch_remaining) { - prefetch_remaining -= bandwidth_idle_time; - bandwidth_idle_time = 0; - VLOG(4) << "Prefetch #" - << (prefetch_completed_index % prefetches.size()) << " (" - << prefetch_completed_index << ") still ongoing at " << i - << ", remaining elapsed = " << prefetch_remaining; - break; - } - bandwidth_idle_time -= prefetch_remaining; - prefetch_remaining = 0; - VLOG(4) << "Prefetch #" - << (prefetch_completed_index % prefetches.size()) << " (" - << prefetch_completed_index << ") completed at " << i - << ", bandwidth idle time = " << bandwidth_idle_time; - } - if (bandwidth_idle_time > 0) { - VLOG(4) << "Bandwidth idle time at " << i << " = " - << bandwidth_idle_time; - total_bandwidth_idle_time += bandwidth_idle_time; - } - - // Start new prefetches that are scheduled to start after this - // instruction. - for (; prefetch_start_index < (iteration + 1) * prefetches.size() && - prefetches[prefetch_start_index % prefetches.size()] - .first->copy_start_schedule_after() == i; - ++prefetch_start_index) { - float& prefetch_remaining = - prefetch_remaining_elapsed_times[prefetch_start_index]; - prefetch_remaining = - prefetches[prefetch_start_index % prefetches.size()].second; - VLOG(4) << "Prefetch #" << (prefetch_start_index % prefetches.size()) - << " (" << prefetch_start_index << ") started at " << i - << ", remaining elapsed = " << prefetch_remaining; - } - } - VLOG(3) << "Iteration " << iteration; - VLOG(3) << "Total elapsed: " << total_elapsed - << ", total critical prefetch: " << total_critical_prefetch - << ", total bandwidth idle time: " << total_bandwidth_idle_time; - result = total_elapsed + total_critical_prefetch; - } - return result; -} - -/*static*/ std::string -MemoryBoundLoopOptimizer::LoopValue::AllocationTypeToString( - LoopValue::AllocationType allocation_type) { - switch (allocation_type) { - case AllocationType::kTemporary: - return "temporary"; - case AllocationType::kLoopCarriedDependence: - return "loop-carried dependence"; - case AllocationType::kPinned: - return "pinned"; - case AllocationType::kPrefetch: - return "prefetch"; - default: - CHECK(allocation_type == AllocationType::kUnsupported); - return "unsupported"; - } -} - -std::string MemoryBoundLoopOptimizer::LoopValue::ToString() const { - std::string values_str; - absl::StrAppend(&values_str, "Values:"); - for (const HloValue* hlo_value : hlo_values) { - absl::StrAppend(&values_str, "\n - ", hlo_value->ToShortString()); - } - std::string allocations_str; - if (!allocations.empty()) { - absl::StrAppend(&allocations_str, "Allocations:"); - } - for (const auto& allocation : allocations) { - absl::StrAppend(&allocations_str, "\n - ", allocation->ToString()); - } - std::string chunk_str; - if (chunks.HasValues()) { - absl::StrAppend(&chunk_str, "\n", - "even chunk: ", chunks.even_chunk->ToString()); - absl::StrAppend(&chunk_str, "\n", - "odd chunk: ", chunks.odd_chunk->ToString()); - absl::StrAppend(&chunk_str, "\n", "alternate memory begin idx in loop: ", - alternate_memory_begin_idx_in_loop.value()); - absl::StrAppend(&chunk_str, "\n", "alternate memory end idx in loop: ", - alternate_memory_end_idx_in_loop.value()); - } - return absl::StrCat( - "Size: ", size, " savings: ", savings, - " savings per byte: ", savings_per_byte, - " allocation type: ", AllocationTypeToString(allocation_type), chunk_str, - "\n", values_str, "\n", allocations_str); -} - -bool MemoryBoundLoopOptimizer::LoopValue::IsAllocationTypeSupported() const { - return allocation_type == AllocationType::kTemporary || - allocation_type == AllocationType::kPinned || - allocation_type == AllocationType::kPrefetch; -} - -void MemoryBoundLoopOptimizer::LoopValue::SetChunkPairAndInterval( - EvenOddChunkPair chunk_pair, int64_t begin_idx_in_loop, - int64_t end_idx_in_loop) { - chunks = chunk_pair; - alternate_memory_begin_idx_in_loop = begin_idx_in_loop; - alternate_memory_end_idx_in_loop = end_idx_in_loop; -} - -void MemoryBoundLoopOptimizer::SortLoopValues() { - absl::c_stable_sort(loop_values_, [](const LoopValue& a, const LoopValue& b) { - return a.savings_per_byte > b.savings_per_byte; - }); -} - -void MemoryBoundLoopOptimizer::AllocateLoopValues() { - // This function allocates loop values. - std::vector prefetch_values; - VLOG(3) << "Pre optimization execution time: " << CalculateExecutionTime(); - for (LoopValue& value : loop_values_) { - switch (value.allocation_type) { - case LoopValue::AllocationType::kTemporary: - AllocateTemporary(value); - break; - case LoopValue::AllocationType::kPinned: - if (value.savings > 0) { - AllocatePinned(value); - } - break; - case LoopValue::AllocationType::kPrefetch: - prefetch_values.push_back(&value); - break; - case LoopValue::AllocationType::kLoopCarriedDependence: - case LoopValue::AllocationType::kUnsupported: - VLOG(1) << "Unsupported allocation: " << value.ToString(); - } - } - VLOG(6) << "Heap after allocating temporaries:\n" - << heap_.MemoryUsageToAsciiArt(); - VLOG(3) << "Execution time after allocating temporaries: " - << CalculateExecutionTime(); - AllocatePrefetches(absl::MakeSpan(prefetch_values)); - VLOG(6) << "Heap after allocating prefetches:\n" - << heap_.MemoryUsageToAsciiArt(); - VLOG(3) << "Execution time after allocating prefetches: " - << CalculateExecutionTime(); -} - -void MemoryBoundLoopOptimizer::PostProcess() { - // At the end, ensure that all loop uses have a corresponding Allocation and - // create one in the default memory space if they don't. - for (LoopValue& value : loop_values_) { - absl::flat_hash_set allocated_uses; - for (const auto& allocation : value.allocations) { - for (const HloUse& use : allocation->uses()) { - allocated_uses.insert(use); - } - } - std::vector unallocated_uses; - absl::flat_hash_set use_indices; - for (const auto& [idx, use] : value.loop_uses) { - use_indices.insert(idx); - if (!allocated_uses.contains(use)) { - unallocated_uses.push_back(use); - } - } - for (const auto& [next_iteration_idx, use] : value.next_iteration_uses) { - if (use_indices.contains(next_iteration_idx)) { - continue; - } - HloInstruction* loop_instruction = - hlo_live_range_.flattened_instruction_sequence().instructions().at( - loop_start_ + next_iteration_idx); - HloUse loop_use{loop_instruction, use.operand_number, use.operand_index}; - if (!allocated_uses.contains(loop_use)) { - unallocated_uses.push_back(loop_use); - } - } - if (!unallocated_uses.empty()) { - // TODO(b/281582241): We should find the correct position. For now, we're - // using the defining position on the first HLO value. - value.allocations.push_back(std::make_unique( - value.hlo_values.front()->defining_position(), MemorySpace::kDefault, - std::nullopt, 0, loop_size_)); - for (const HloUse& use : unallocated_uses) { - value.allocations.back()->AddUse(use); - } - } - VLOG(3) << "LoopValue: " << value.ToString(); - } -} - -bool MemoryBoundLoopOptimizer::AllocateTemporary(LoopValue& value) { - VLOG(3) << "AllocateTemporary: " << value.ToString(); - if (value.hlo_values.size() > 1) { - VLOG(3) << "LoopValue has more than one hlo value associated."; - return false; - } - int64_t definition_idx = value.loop_positions.front().first; - int64_t max_use_idx; - int64_t begin_idx_in_loop = definition_idx; - int64_t end_idx_in_loop; - if (!value.next_iteration_uses.empty()) { - max_use_idx = value.next_iteration_uses.back().first; - // If max_use_idx >= definition_idx, then this is a loop carried dependence - // and we should not have called this function. - CHECK_LT(max_use_idx, definition_idx); - end_idx_in_loop = max_use_idx + loop_size_; - } else { - max_use_idx = value.loop_uses.back().first; - end_idx_in_loop = max_use_idx; - } - EvenOddChunkPair chunks = heap_.AllocateSameEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, value.size); - if (!chunks.HasValues()) { - VLOG(3) << "Could not find Allocation for temporary value: " - << value.ToString(); - return false; - } - value.SetChunkPairAndInterval(chunks, begin_idx_in_loop, end_idx_in_loop); - VLOG(3) << "Pos: " << value.loop_positions[0].second; - VLOG(3) << "Allocation found for temporary value: " << value.ToString(); - VLOG(6) << "Heap after allocating temporary value: " - << heap_.MemoryUsageToAsciiArt(); - value.allocations.push_back(std::make_unique( - value.loop_positions[0].second, MemorySpace::kAlternate, std::nullopt, - definition_idx, max_use_idx)); - AddAllLoopPositionsAndUses(value, /*allocate_next_iteration_uses=*/true); - return true; -} - -bool MemoryBoundLoopOptimizer::AllocatePinned(LoopValue& value) { - int64_t begin_idx_in_loop = 0; - int64_t end_idx_in_loop = loop_size_ - 1; - EvenOddChunkPair chunks = heap_.AllocateSameEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, value.size); - if (!chunks.HasValues()) { - VLOG(3) << "Could not find Allocation for pinned value: " - << value.ToString(); - return false; - } - value.SetChunkPairAndInterval(chunks, begin_idx_in_loop, end_idx_in_loop); - CHECK(value.header_position); - VLOG(3) << "Allocation found for pinned value: " << value.ToString(); - VLOG(6) << "Heap after allocating pinned value: " - << heap_.MemoryUsageToAsciiArt(); - value.allocations.push_back(std::make_unique( - *value.header_position, MemorySpace::kAlternate, std::nullopt, 0, - loop_size_)); - AddAllLoopPositionsAndUses(value, /*allocate_next_iteration_uses=*/false); - return true; -} - -bool MemoryBoundLoopOptimizer::AllocatePrefetches( - absl::Span values) { - VLOG(3) << "Allocating prefetches num values: " << values.size(); - AllocatePrefetchesContext context; - context.values = values; - // Populate value_indices, which is a list of indices into values array sorted - // by the start time of the first use. - context.value_indices.resize(values.size()); - absl::c_iota(context.value_indices, 0); - absl::c_stable_sort(context.value_indices, [&](int a, int b) { - return std::forward_as_tuple( - values[a]->loop_uses.begin()->first, - values[a]->loop_uses.begin()->second.operand_number) > - std::forward_as_tuple( - values[b]->loop_uses.begin()->first, - values[b]->loop_uses.begin()->second.operand_number); - }); - - // Populate the data structures that contain additional positions and uses - // that would get alternate memory allocations if all of the prefetches were - // successful. - absl::flat_hash_map>> - additional_uses_in_alternate_mem; - absl::flat_hash_map> - additional_positions_in_alternate_mem; - for (const LoopValue* value : values) { - VLOG(3) << " prefetch value: " << value->ToString(); - for (const auto& [idx, use] : value->loop_uses) { - additional_uses_in_alternate_mem[use.instruction].push_back( - {use.operand_number, use.operand_index}); - } - for (const auto& [idx, position] : value->loop_positions) { - additional_positions_in_alternate_mem[position.instruction].push_back( - position.index); - } - } - // Calculate the default-memory remaining bandwidths assuming all prefetches - // succeed. - for (int i = 0; i < loop_size_; ++i) { - context.bandwidth_idle_times.push_back( - GetBandwidthIdleTime(i, additional_uses_in_alternate_mem, - additional_positions_in_alternate_mem)); - VLOG(3) << "Remaining bandwidth at " << i << " = " - << *context.bandwidth_idle_times.rbegin(); - } - - // Allocate prefetches by traversing the loop values in reverse order of - // the first uses. - for (int value_index : context.value_indices) { - AllocatePrefetch(value_index, context); - } - - for (int i = 0; i < loop_size_; ++i) { - VLOG(3) << "Remaining bandwidth [" << i - << "] : " << context.bandwidth_idle_times[i]; - } - return true; -} - -bool MemoryBoundLoopOptimizer::AllocatePrefetch( - int value_index, AllocatePrefetchesContext& context) { - LoopValue* value = context.values[value_index]; - VLOG(3) << "Allocating value: " << value->ToString(); - int first_use_idx = value->loop_uses.front().first; - int last_use_idx = value->loop_uses.back().first; - int last_use_idx_sentinel = last_use_idx; - if (!value->next_iteration_uses.empty()) { - last_use_idx = value->next_iteration_uses.back().first; - last_use_idx_sentinel = last_use_idx + loop_size_; - CHECK_LT(last_use_idx, first_use_idx); - } - float copy_resource = cost_analysis_.GetAsyncCopyElapsed( - cost_analysis_.GetShapeSizeBytes(value->hlo_values.front()->shape())); - VLOG(3) << "First use: " << value->loop_uses.begin()->second - << " use idx: " << first_use_idx - << " copy resource: " << copy_resource; - const auto& [even_chunk, odd_chunk] = heap_.FindEvenAndOddAllocationBetween( - first_use_idx, last_use_idx_sentinel, value->size); - if (!even_chunk.has_value() || !odd_chunk.has_value()) { - // Not enough memory to even fit the value in the alternate memory for the - // duration of its live range. - VLOG(3) << "Could not find Allocation for prefetch value: " - << value->ToString(); - return false; - } - - std::optional copy_start_loop_idx; - int committed_early_forced_prefetches_count = 0; - // The general allocation algorithm for prefetches is to first calculate the - // default-memory bandwidth idle times at each point (assuming all prefetches - // succeeded). We show this pictorially below. We also show the previous - // iteration for clarity. The algorithm solves allocation for one iteration - // and this will be used for all iterations. - // - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 2 2 1 2 3 1| 2 2 1 2 3 1| - // additional memory: 0 0 0 0 0 0| 0 0 0 0 0 0| - // iteration: prev | current | - // - // Now, let's assume there are two prefetches that need to be scheduled. For - // the sake of the example, assume 1 MiB of prefetch uses 1 memory bandwidth - // resource: - // - Prefetch 1 is 4 MiB and is first used at index 5. - // - Prefetch 2 is 5 MiB and is first used at index 1. - // - // We first order these prefetches by their first use from latest to earliest. - // Then starting from the prefetch completion time (i.e. the first use time), - // move the prefetch start time earlier until the copy resource is satisfied - // (or reaching another resource satisfaction criteria explained below) by - // consuming the bandwidth idle time of the overlapped instructions. We also - // keep track of the additional memory required. Note that index 5 also - // accounts for the additional 4 MiB consumed since the data needs to reside - // during the execution of the instruction at index 5. Below is the updated - // state after scheduling prefetch 1: - // - // prefetch 1: +====+ +====+ - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 2 2 1 1 0 1| 2 2 1 1 0 1| - // additional memory: 0 0 0 4 4 4| 0 0 0 4 4 4| - // iteration: prev | current | - // - // To schedule prefetch 2, we similarly start the same way, from its first use - // and bring the prefetch start earlier. We first reach index 0 with still an - // unsatisfied copy resource of 3: - // - // prefetch 2: +=+ +=+ unsat res: 3 - // prefetch 1: +====+ +====+ - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 0 2 1 1 0 1| 0 2 1 1 0 1| - // additional memory: 5 5 0 4 4 4| 5 5 0 4 4 4| - // iteration: prev | current | - // - // We continue onto the previous iteration: - // - // prefetch 2:===+ +====+ +== unsat res: 2 - // prefetch 1: +====+ +====+ - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 0 2 1 1 0 0| 0 2 1 1 0 0| - // additional memory: 5 5 0 4 4 9| 5 5 0 4 4 9| - // iteration: prev | current | - // - // As we bring the start time of prefetch 2 earlier, it starts overlapping - // with prefetch 1: - // - // prefetch 2:===+ +==========+ +======== unsat res: 1 - // prefetch 1: +====+ +====+ - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 0 2 1 0 0 0| 0 2 1 0 0 0| - // additional memory: 5 5 0 9 9 9| 5 5 0 9 9 9| - // iteration: prev | current | - // - // The prefetch resource is still unsatisfied at this point. We can bring the - // prefetch earlier. However, the first prefetch's end time is earlier than - // the second and we need to maintain FIFO order with regard to prefetches. In - // order to maintain this FIFO order, we "early force" prefetches that are - // already scheduled by moving the start time earlier along with prefetch 2: - // - // prefetch 2:===+ +=============+ +=========== - // prefetch 1: +=======+ +=======+ - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 0 2 0 0 0 0| 0 2 0 0 0 0| - // additional memory: 5 5 9 9 9 9| 5 5 9 9 9 9| - // iteration: prev | current | - // - // Depending on the options provided, we can use alternative resource - // satisfaction criteria. One option is to specify a percentage of the copy - // resource that needs to be satisfied instead of the complete amount (100%). - // This is called the "desired copy ratio". The reason why desired copy ratio - // can be less than 100% is that in a memory-bound loop, we probably do not - // have enough aggregate bandwidth resources to satisfy all of the prefetches, - // but using up all of the default-memory bandwidth is more important than - // having some prefetches with unsatisfied resources. In a similar vein, - // another option is to accept prefetches that are fully pipelined, i.e. - // their copy start time is scheduled the same time as the copy done time in - // the previous iteration, regardless of how much of its copy resources are - // actually satisfied. To illustrate a fully pipelined prefetch, consider - // prefetch 3 (assume no prefetch 1 or 2 in this example) which is 15 MiB and - // its first use is at index 4: - // - // prefetch 3:=============+=================+===== unsat res: 4 - // idx: 0 1 2 3 4 5| 0 1 2 3 4 5| - // bw idle time: 0 0 0 0 0 0| 0 0 0 0 0 0| - // additional memory: 15 15 15 15 30 15|15 15 15 15 30 15| - // iteration: prev | current | - // - // Note that the additional memory consumption at index 4 is actually twice - // the size of the prefetch as we are effectively double buffering. Also note - // that the prefetch has an unsatisfied copy resource of 4 meaning the copy - // will be in the critical path, but this actually will be faster than not - // scheduling this particular prefetch in the first place since the bandwidth - // idle time resource would go unused. - float accumulated_copy_resource = 0; - std::vector early_forced_prefetch_value_indices; - int early_forced_prefetch_value_search_index = 0; - VLOG(6) << "Memory usage before allocating prefetch value: " - << value->ToString() << "\n" - << heap_.MemoryUsageToAsciiArt(); - // NOTE: We can, in practice, run the following loop for loop_size - // iterations(one full loop), till first_use_idx - loop_size, as opposed to - // limiting it till last_use_idx_sentinel - loop_size. This will allow a - // prefetch to use all the idle bandwidth available during one full loop - // iteration. - for (int current_idx = first_use_idx - 1; - current_idx >= last_use_idx_sentinel - loop_size_; --current_idx) { - int loop_idx = (current_idx + loop_size_) % loop_size_; - // Check if this prefetch rolls over to the previous iteration, check if any - // already-scheduled prefetches would violate the FIFO order, and if so, - // "early-force" them to be co-scheduled with this prefetch to maintain the - // FIFO order. This of course increases the required memory, so also keep - // track of additional memory that would be consumed. - if (current_idx < 0) { - for (; context.value_indices[early_forced_prefetch_value_search_index] != - value_index; - ++early_forced_prefetch_value_search_index) { - VLOG(3) << "Searching for early forced: " - << early_forced_prefetch_value_search_index; - LoopValue* early_forced_value = - context.values[context.value_indices - [early_forced_prefetch_value_search_index]]; - if (early_forced_value->allocations.empty()) { - continue; - } - const CopyAllocation* early_forced_prefetch = - static_cast( - early_forced_value->allocations.back().get()); - VLOG(3) << "Prefetch: " << early_forced_prefetch->ToString(); - - // If the prefetch is already a roll-around prefetch, no need to further - // early force it. - if (early_forced_prefetch->copy_done_schedule_before() <= - early_forced_prefetch->copy_start_schedule_after() + 1 || - (early_forced_prefetch->copy_start_schedule_after() == - loop_size_ - 1 && - early_forced_prefetch->copy_done_schedule_before() == 0)) { - break; - } - if (early_forced_prefetch->copy_start_schedule_after() != loop_idx) { - break; - } - early_forced_prefetch_value_indices.push_back( - early_forced_prefetch_value_search_index); - VLOG(6) - << "Memory usage before removing prefetch value for early force: " - << early_forced_value->ToString() << "\n" - << heap_.MemoryUsageToAsciiArt(); - // Remove the original chunk from the heap. - heap_.RemoveEvenOddChunkPair( - early_forced_value->alternate_memory_begin_idx_in_loop.value(), - early_forced_value->alternate_memory_end_idx_in_loop.value(), - early_forced_value->chunks); - } - } - - VLOG(3) << "Loop idx:" << loop_idx << " Early force prefetch values: " - << early_forced_prefetch_value_indices.size(); - VLOG(6) << "Memory usage before adding pending chunks: \n" - << heap_.MemoryUsageToAsciiArt(); - std::vector pending_chunk_intervals; - for (int early_forced_prefetch_value_index : - early_forced_prefetch_value_indices) { - LoopValue* early_forced_value = - context - .values[context.value_indices[early_forced_prefetch_value_index]]; - int64_t begin_idx_in_loop = loop_idx; - int64_t end_idx_in_loop = - early_forced_value->alternate_memory_end_idx_in_loop.value(); - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, early_forced_value->size); - if (!chunks.HasValues()) { - VLOG(3) << "Could not allocate between " << begin_idx_in_loop << " and " - << end_idx_in_loop << " for early forced value: " - << early_forced_value->ToString(); - VLOG(6) << "Memory usage after failed allocation: \n" - << heap_.MemoryUsageToAsciiArt(); - break; - } - pending_chunk_intervals.push_back( - {begin_idx_in_loop, end_idx_in_loop, chunks}); - VLOG(3) << "Added pending chunk: " - << pending_chunk_intervals.back().ToString() - << " for value: " << early_forced_value->ToString(); - } - - if (pending_chunk_intervals.size() == - early_forced_prefetch_value_indices.size()) { - int64_t begin_idx_in_loop = current_idx; - int64_t end_idx_in_loop = last_use_idx_sentinel; - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, value->size); - if (chunks.HasValues()) { - pending_chunk_intervals.push_back( - {begin_idx_in_loop, end_idx_in_loop, chunks}); - VLOG(3) << "Added pending chunk: " - << pending_chunk_intervals.back().ToString() - << " for current value: " << value->ToString(); - } else { - VLOG(3) << "Could not allocate between " << begin_idx_in_loop << " and " - << end_idx_in_loop << " for value: " << value->ToString(); - VLOG(6) << "Memory usage after failed allocation: \n" - << heap_.MemoryUsageToAsciiArt(); - } - } - - bool out_of_memory = pending_chunk_intervals.size() < - early_forced_prefetch_value_indices.size() + 1; - - // Remove the pending chunks from the heap. - for (auto& pending_chunk_interval : pending_chunk_intervals) { - VLOG(3) << "Removing pending chunk: " - << pending_chunk_interval.ToString(); - heap_.RemoveEvenOddChunkPair(pending_chunk_interval.begin_idx_in_loop, - pending_chunk_interval.end_idx_in_loop, - pending_chunk_interval.chunks); - } - - VLOG(6) << "Memory usage after removing pending chunks: " - << heap_.MemoryUsageToAsciiArt(); - - if (out_of_memory) { - VLOG(3) << "Ran out of memory for value: " << value->ToString(); - break; - } - - // We ideally find a time to overlap the prefetch fully where the previous - // iteration's memory use is disjoint from this iteration. If that is not - // possible, there are two compromises we could pick: - // - Find a prefetch time that satisfies a desired ratio < 1 of the - // prefetch elapsed time. This means the prefetch will be critical. - // - Overlap the prefetch with the previous iteration's buffer use, i.e. - // full pipelining. This would increase the peak memory consumption. - float bandwidth_idle_time = context.bandwidth_idle_times[loop_idx]; - VLOG(3) << "Idx " << loop_idx - << " bandwidth_idle_time: " << bandwidth_idle_time - << " copy resource remaining: " - << (copy_resource - accumulated_copy_resource) << " diff: " - << (bandwidth_idle_time - - (copy_resource - accumulated_copy_resource)); - if (bandwidth_idle_time >= copy_resource - accumulated_copy_resource) { - accumulated_copy_resource = copy_resource; - copy_start_loop_idx = current_idx; - committed_early_forced_prefetches_count = - early_forced_prefetch_value_indices.size(); - VLOG(3) << "Found the complete copy ratio and updated accumulated copy " - "resource: " - << accumulated_copy_resource; - break; - } else if (!copy_start_loop_idx.has_value() && - accumulated_copy_resource + bandwidth_idle_time >= - copy_resource * options_.desired_copy_ratio()) { - accumulated_copy_resource += bandwidth_idle_time; - copy_start_loop_idx = current_idx; - committed_early_forced_prefetches_count = - early_forced_prefetch_value_indices.size(); - VLOG(3) << "Found the desired copy ratio and updated accumulated copy " - "resource: " - << accumulated_copy_resource; - } else if (options_.allow_unsatisfied_fully_pipelined_prefetch() && - loop_idx == last_use_idx) { - // Even if desired resource isn't reached, and if the options allow it, - // allow a fully pipelined prefetch. - accumulated_copy_resource += bandwidth_idle_time; - copy_start_loop_idx = current_idx; - committed_early_forced_prefetches_count = - early_forced_prefetch_value_indices.size(); - VLOG(3) << "Could not reach the desired copy ratio but scheduling " - "fully pipelined prefetch anyway: " - << accumulated_copy_resource; - break; - } else { - accumulated_copy_resource += bandwidth_idle_time; - VLOG(3) << "Updated accumulated copy resource: " - << accumulated_copy_resource; - } - } - - // Restore original heap state as is for values that are not being early - // forced. This is to ensure that the memory usage is the same as before early - // forcing. If no copy start time was found, all the prefetches will be - // restored to their original state. If a copy start time was found, the - // prefetches that will not be early forced will be restored to their original - // state. - VLOG(6) << "Memory usage before restoring original state: " - << heap_.MemoryUsageToAsciiArt(); - for (int i = committed_early_forced_prefetches_count; - i < early_forced_prefetch_value_indices.size(); ++i) { - int early_forced_prefetch_value_index = - early_forced_prefetch_value_indices[i]; - LoopValue* early_forced_value = - context - .values[context.value_indices[early_forced_prefetch_value_index]]; - // Allocate a chunk in at the same offset as the original prefetch. - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - early_forced_value->alternate_memory_begin_idx_in_loop.value(), - early_forced_value->alternate_memory_end_idx_in_loop.value(), - early_forced_value->size, - {early_forced_value->chunks.even_chunk->offset, - early_forced_value->chunks.odd_chunk->offset}); - // The chunk should always be present as we are allocating at the same - // offset. - CHECK(chunks.HasValues()); - CHECK_EQ(chunks.even_chunk->offset, - early_forced_value->chunks.even_chunk->offset); - CHECK_EQ(chunks.odd_chunk->offset, - early_forced_value->chunks.odd_chunk->offset); - } - VLOG(6) << "Memory usage after restoring original state: " - << heap_.MemoryUsageToAsciiArt(); - - if (!copy_start_loop_idx.has_value()) { - VLOG(3) << "Could not find a suitable copy start time for value: " - << value->ToString(); - return false; - } - - VLOG(3) << "Success: copy_start_loop_idx: " << copy_start_loop_idx.value() - << " leftover copy resource: " - << (copy_resource - accumulated_copy_resource); - // We are early forcing the prefetches of the previous iteration. This is the - // corresponding copy start index in the previous iteration. - int early_prefetch_copy_start_loop_idx = - (copy_start_loop_idx.value() + loop_size_) % loop_size_; - // We reset accumulated copy resource and then reuse it to accumulate copy - // resource time in order to replay the previous for loop. It is important - // that we use the same arithmetic operations (as opposed to subtracting from - // copy_resource) because floating point operations aren't commutative. - accumulated_copy_resource = 0.0; - for (int i = first_use_idx - 1; i >= last_use_idx_sentinel - loop_size_; - --i) { - int loop_idx = (i + loop_size_) % loop_size_; - float& bandwidth_idle_time = context.bandwidth_idle_times[loop_idx]; - if (bandwidth_idle_time < copy_resource - accumulated_copy_resource) { - accumulated_copy_resource += bandwidth_idle_time; - bandwidth_idle_time = 0; - if (loop_idx == early_prefetch_copy_start_loop_idx) { - VLOG(3) << "Remaining copy resource: " - << (copy_resource - accumulated_copy_resource); - break; - } - } else { - bandwidth_idle_time -= copy_resource - accumulated_copy_resource; - CHECK_EQ(loop_idx, early_prefetch_copy_start_loop_idx); - break; - } - } - - // Account for the additional memory used by early forcing the already - // scheduled prefetches. Also modify the start times of these to this - // prefetch's copy start time. - // Allocate the force-early prefetches first, and allocate them in the same - // order as we did to check for out-of-memory, so we can reproduce the same - // allocation pattern. - // TODO(subhankarshah): Instead of depending on the order of allocation, store - // the offsets of the early forced prefetches and use that to allocate them. - for (int i = 0; i < committed_early_forced_prefetches_count; ++i) { - int early_forced_prefetch_value_index = - early_forced_prefetch_value_indices[i]; - LoopValue* early_forced_value = - context - .values[context.value_indices[early_forced_prefetch_value_index]]; - CHECK(!early_forced_value->allocations.empty()); - CopyAllocation* early_forced_prefetch = static_cast( - early_forced_value->allocations.back().get()); - int64_t begin_idx_in_loop = early_prefetch_copy_start_loop_idx; - int64_t end_idx_in_loop = - early_forced_value->alternate_memory_end_idx_in_loop.value(); - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, early_forced_value->size); - // The chunk should always be present as we reproducing the same allocation - // pattern as the out-of-memory check. - CHECK(chunks.HasValues()); - CHECK_LT(begin_idx_in_loop, - early_forced_value->alternate_memory_begin_idx_in_loop.value()); - early_forced_value->SetChunkPairAndInterval(chunks, begin_idx_in_loop, - end_idx_in_loop); - early_forced_prefetch->set_copy_start_schedule_after( - ((early_prefetch_copy_start_loop_idx - 1) + loop_size_) % loop_size_); - VLOG(3) << "Early forced prefetch: " << early_forced_value->ToString(); - VLOG(6) << "Memory usage after allocating early forced prefetch: " - << heap_.MemoryUsageToAsciiArt(); - } - - // Create the Allocation objects that correspond to the scheduled prefetch. - CHECK(value->header_position); - value->allocations.push_back(std::make_unique( - *value->header_position, MemorySpace::kDefault, std::nullopt, 0, - loop_size_)); - int64_t begin_idx_in_loop = copy_start_loop_idx.value(); - int64_t end_idx_in_loop = last_use_idx_sentinel; - // The chunk should always be present as we reproducing the same allocation - // pattern as the out-of-memory check. - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, value->size); - CHECK(chunks.HasValues()); - value->SetChunkPairAndInterval(chunks, begin_idx_in_loop, end_idx_in_loop); - value->allocations.push_back(std::make_unique( - *value->allocations.back(), MemorySpace::kAlternate, std::nullopt, - ((early_prefetch_copy_start_loop_idx - 1) + loop_size_) % loop_size_, - first_use_idx, last_use_idx_sentinel)); - VLOG(3) << "Allocation found for prefetch: " << value->ToString(); - VLOG(6) << "Memory usage after allocating prefetch: " << value->ToString() - << "\n" - << heap_.MemoryUsageToAsciiArt(); - AddAllLoopPositionsAndUses(*value, /*allocate_next_iteration_uses=*/true); - return true; -} - -void MemoryBoundLoopOptimizer::AddAllLoopPositionsAndUses( - LoopValue& value, bool allocate_next_iteration_uses) { - CHECK_GE(value.allocations.size(), 1); - Allocation& allocation = *value.allocations.back(); - for (const auto& [idx, position] : value.loop_positions) { - positions_in_alternate_mem_[position.instruction].push_back(position.index); - } - for (const auto& [idx, use] : value.loop_uses) { - uses_in_alternate_mem_[use.instruction].push_back( - {use.operand_number, use.operand_index}); - allocation.AddUse(use); - } - if (allocate_next_iteration_uses) { - for (const auto& [next_iteration_idx, use] : value.next_iteration_uses) { - HloInstruction* loop_instruction = - hlo_live_range_.flattened_instruction_sequence().instructions().at( - loop_start_ + next_iteration_idx); - uses_in_alternate_mem_[loop_instruction].push_back( - {use.operand_number, use.operand_index}); - allocation.AddUse( - {loop_instruction, use.operand_number, use.operand_index}); - } - } -} - -float MemoryBoundLoopOptimizer::GetBandwidthIdleTime(int idx) const { - const HloInstruction* inst = - hlo_live_range_.flattened_instruction_sequence().instructions().at( - loop_start_ + idx); - std::vector> empty_operands; - std::vector empty_outputs; - const std::vector>* operands_in_alternate_mem = - &empty_operands; - const std::vector* outputs_in_alternate_mem = &empty_outputs; - auto uses_it = uses_in_alternate_mem_.find(inst); - if (uses_it != uses_in_alternate_mem_.end()) { - operands_in_alternate_mem = &uses_it->second; - } - auto positions_it = positions_in_alternate_mem_.find(inst); - if (positions_it != positions_in_alternate_mem_.end()) { - outputs_in_alternate_mem = &positions_it->second; - } - return cost_analysis_.GetDefaultMemoryBandwidthIdleTime( - *inst, *operands_in_alternate_mem, *outputs_in_alternate_mem); -} - -float MemoryBoundLoopOptimizer::GetBandwidthIdleTime( - int idx, - const absl::flat_hash_map>>& - additional_uses_in_alternate_mem, - const absl::flat_hash_map>& - additional_positions_in_alternate_mem) const { - const HloInstruction* inst = - hlo_live_range_.flattened_instruction_sequence().instructions().at( - loop_start_ + idx); - std::vector> operands_in_alternate_mem; - std::vector outputs_in_alternate_mem; - auto uses_it = uses_in_alternate_mem_.find(inst); - if (uses_it != uses_in_alternate_mem_.end()) { - operands_in_alternate_mem = uses_it->second; - } - auto additional_uses_it = additional_uses_in_alternate_mem.find(inst); - if (additional_uses_it != additional_uses_in_alternate_mem.end()) { - absl::c_copy(additional_uses_it->second, - std::back_inserter(operands_in_alternate_mem)); - } - auto positions_it = positions_in_alternate_mem_.find(inst); - if (positions_it != positions_in_alternate_mem_.end()) { - outputs_in_alternate_mem = positions_it->second; - } - auto additional_positions_it = - additional_positions_in_alternate_mem.find(inst); - if (additional_positions_it != additional_positions_in_alternate_mem.end()) { - absl::c_copy(additional_positions_it->second, - std::back_inserter(outputs_in_alternate_mem)); - } - return cost_analysis_.GetDefaultMemoryBandwidthIdleTime( - *inst, operands_in_alternate_mem, outputs_in_alternate_mem); -} - -float MemoryBoundLoopOptimizer::GetInstructionElapsed(int idx) const { - const HloInstruction* inst = - hlo_live_range_.flattened_instruction_sequence().instructions().at( - loop_start_ + idx); - std::vector> empty_operands; - std::vector empty_outputs; - const std::vector>* operands_in_alternate_mem = - &empty_operands; - const std::vector* outputs_in_alternate_mem = &empty_outputs; - auto uses_it = uses_in_alternate_mem_.find(inst); - if (uses_it != uses_in_alternate_mem_.end()) { - operands_in_alternate_mem = &uses_it->second; - } - auto positions_it = positions_in_alternate_mem_.find(inst); - if (positions_it != positions_in_alternate_mem_.end()) { - outputs_in_alternate_mem = &positions_it->second; - } - return cost_analysis_.GetInstructionElapsedInAlternateMemory( - *inst, *operands_in_alternate_mem, *outputs_in_alternate_mem); -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h deleted file mode 100644 index a5d20c5a70b6de..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h +++ /dev/null @@ -1,484 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_MEMORY_BOUND_LOOP_OPTIMIZER_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_MEMORY_BOUND_LOOP_OPTIMIZER_H_ - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/base/nullability.h" -#include "absl/container/flat_hash_map.h" -#include "absl/log/check.h" -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/buffer_value.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo.pb.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -// Pair of chunks for even and odd loop iterations. -struct EvenOddChunkPair { - std::optional even_chunk; - std::optional odd_chunk; - - bool HasValues() const { - return even_chunk.has_value() && odd_chunk.has_value(); - } -}; - -// LoopOptimizerBestFitHeap extends GlobalDecreasingSizeBestFitHeap to track -// allocated buffers and their live intervals for the MemoryBoundLoopOptimizer. -// * We model 4 loop iterations. -// * The 0th and 2nd iterations are even. The 1st and 3rd are odd. -// * Allocations in even iterations are required to have the same offsets. -// Likewise, allocations in odd iterations are required to have the same -// offset. -// * Allocations may have different offsets between odd and even iterations. -// * Buffers can span up to 2 iterations. -// * The algorithm uses the 0th and 1st iterations to account for buffers that -// start in those iterations but are still alive in the 2nd and 3rd -// iterations. The 2nd and 3rd iterations are used to give the complete loop -// buffer picture. -class LoopOptimizerBestFitHeap - : public GlobalDecreasingSizeBestFitHeap { - public: - explicit LoopOptimizerBestFitHeap(uint64_t size_limit_per_heap, - int64_t loop_size, - int64_t alignment_in_bytes) - : GlobalDecreasingSizeBestFitHeap(alignment_in_bytes), - size_limit_per_heap_(size_limit_per_heap), - loop_size_(loop_size) {} - ~LoopOptimizerBestFitHeap() override = default; - - // Frees the memory space denoted by chunk from [begin_idx_in_loop, - // end_idx_in_loop] from all iterations. - void RemoveEvenOddChunkPair(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - EvenOddChunkPair& chunks); - - // Displays the current memory usage vs time for 6 loop iterations by default. - // Note: The 0th and the 1st iterations are just to account for loop around - // for buffers that go across one or two loop boundaries. The 2nd and the 3rd - // iterations present the actual memory view of the allocation. The 4th and - // 5th iterations show buffers from previous two iterations that go across one - // or two loop boundaries. begin_iteration_idx and end_iteration_idx are both - // inclusive, 0 indexed. - std::string MemoryUsageToAsciiArt(int64_t begin_iteration = 0, - int64_t end_iteration = 5) const; - - // Returns a vector of size loop_size, where the i'th element denotes the - // available(unfragmented) alternate memory in bytes at loop_idx i. - std::vector RemainingMemoryByTime() const; - - // Returns an integer denoting the largest occupied memory location in the - // alternate memory. - int64_t LastMemoryOffsetOccupied() const; - - // Finds free memory chunks of size "size" between [begin_idx_in_loop, - // end_idx_in_loop] in the even and odd loop iterations, only if free chunks - // are found in both iterations. The even and odd iteration offsets may be - // different. - EvenOddChunkPair FindEvenAndOddAllocationBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - std::pair preferred_offsets = {-1, -1}); - - // Finds and reserves free memory chunks of size "size" between - // [begin_idx_in_loop, end_idx_in_loop] in the even and odd loop iterations, - // only if free chunks are found in both iterations. The even and odd - // iteration offsets may be different. - EvenOddChunkPair AllocateEvenAndOddBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - std::pair preferred_offsets = {-1, -1}); - - // Finds free memory chunks of size "size" between [begin_idx_in_loop, - // end_idx_in_loop] in the even and odd loop iterations, only if free chunks - // are found in both iterations. The even and odd iteration offsets are same. - EvenOddChunkPair FindSameEvenAndOddAllocationBetween( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size, - int64_t preferred_offset = -1); - - // Finds and reserves free memory chunks of size "size" between - // [begin_idx_in_loop, end_idx_in_loop] in the even and odd loop iterations, - // only if free chunks are found in both iterations. The even and odd - // iteration offsets are same. - EvenOddChunkPair AllocateSameEvenAndOddBetween(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size, - int64_t preferred_offset = -1); - - private: - // REQUIRES: - // * begin_idx_in_loop <= end_idx_in_loop - // * begin_idx_in_loop is within [-loop_size loop_size) - // * end_idx_in_loop is within [0, 2 * loop_size) - // * end_idx_in_loop - begin_idx_in_loop + 1 <= 2 * loop_size (allocation - // colocated in even (or odd) iterations cannot span more than 2 loop - // iterations) - void CheckAllocationIntervalValid(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop) const; - - // Shifts allocation interval at [begin_idx_in_loop, end_idx_in_loop] to - // [begin_idx_in_loop + loop_size, end_idx_in_loop + loop_size], if - // begin_idx_in_loop is negative. - void ShiftAllocationIntervalIfRequired(int64_t& begin_idx_in_loop, - int64_t& end_idx_in_loop) const; - - // Returns pointer to a newly created allocation block that is added to - // allocation_blocks_. - const AllocationBlock& GetAllocationBlock(int64_t start_time, - int64_t end_time, int64_t size); - - // Creates a BufferInterval corresponding to the AllocationBlock and adds it - // to buffer_intervals_. - void CreateBufferInterval(const AllocationBlock& allocation_block, - const AllocationBlock* colocated_with = nullptr); - - std::optional MaybeFindChunkCandidate( - const AllocationBlock& allocation_block, int64_t preferred_offset = -1); - - std::optional FindAndCommitChunkCandidate( - const AllocationBlock& allocation_block, int64_t preferred_offset = -1); - - void RemoveChunk(int64_t start_time, int64_t end_time, Chunk chunk); - - void RemoveEvenChunks(int64_t begin_idx_in_loop, int64_t end_idx_in_loop, - std::optional& chunk); - - void RemoveOddChunks(int64_t begin_idx_in_loop, int64_t end_idx_in_loop, - std::optional& chunk); - - // Creates colocated allocation blocks for loop interval [begin_idx_in_loop, - // end_idx_in_loop] in 0th, 1st, 2nd and 3rd loop iterations and returns one - // AllocationBlock. - const AllocationBlock& CreateSameEvenAndOddAllocationBlock( - int64_t begin_idx_in_loop, int64_t end_idx_in_loop, int64_t size); - - // Creates colocated allocation blocks for loop interval [begin_idx_in_loop, - // end_idx_in_loop] in 0th and 2nd loop iterations and returns one - // AllocationBlock. - const AllocationBlock& CreateEvenAllocationBlock(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size); - // Creates colocated allocation blocks for loop interval [begin_idx_in_loop, - // end_idx_in_loop] in 1st and 3rd loop iterations and returns one - // AllocationBlock. - const AllocationBlock& CreateOddAllocationBlock(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size); - - uint64_t size_limit_per_heap_; - int64_t loop_size_; - std::list allocation_blocks_; -}; - -// An optimizer for unrolled memory-bound loops. It keeps track of alternate -// memory capacity and default memory bandwidth to decide the allocations of -// each tensor within a loop iteration. The assumption is that all of the -// unrolled loop iterations will use the same allocation decisions, so we can -// spend more time to optimize this one iteration as optimally as possible. -// -// To represent instructions, we keep track of three iterations (previous, -// current, and next), as well as the header and footer regions that are before -// and after the loop, respectively. -// -// We classify each tensor used in the current iteration as one of the following -// allocations based on its positions and uses: -// -// Temporary Allocations: These are produced by a producer in the current -// iteration and consumed either in this or the next iteration. For these, we -// try to give them alternate memory allocations for their entire live range. -// -// Case 1: producer and consumer all in the current iteration. -// p-----c--c -// Case 2: producer is in the current iter, consumer is in the next iter. -// p-----c -// idx: |...| 0 1 2 3 4| 0 1 2 3 4| 0 1 2 3 4|...| -// iter: head |...| prev | current | next |...| foot -// -// Loop Carried Dependences: This is where the last use is at a larger index -// than the producer. This would require 2X peak buffer consumption because both -// this and next iteration's buffer is alive at the same time. This case is -// currently not supported. -// -// Case 3: producer is in the current iter, consumer is in the next iter -// (consumer idx >= producer idx). -// p-----------------c -// idx: |...| 0 1 2 3 4| 0 1 2 3 4| 0 1 2 3 4|...| -// iter: head |...| prev | current | next |...| foot -// -// Pinned Allocations: These are values produced at the header and are used in -// every iteration at the same indices. For these, we just allocate the buffer -// for the duration of the loop: -// -// Case 4: producer: kHead, consumer: kCurrent -// p---------------c--------------c--------------c-------- -// idx: |...| 0 1 2 3 4| 0 1 2 3 4| 0 1 2 3 4|...| -// iter: head |...| prev | current | next |...| foot -// -// Prefetch Allocations: These are values produced at the header and are used in -// the current (and possibly next) iteration. We will try to prefetch these -// values into the alternate memory: -// -// Case 5: producer: kHead, consumer: kCurrent -// p---------------------------------c--------c -// idx: |...| 0 1 2 3 4| 0 1 2 3 4| 0 1 2 3 4|...| -// iter: head |...| prev | current | next |...| foot -class MemoryBoundLoopOptimizer { - public: - // We represent each tensor used in the current iteration as a LoopValue, - // wrapping the relevant information such as its HLO value, indices and - // pointers to its use and position sites in different iterations. - // TODO(b/364621066): Make LoopValue a class. - struct LoopValue { - // An enum that encodes the allocation type that is suitable for this - // LoopValue. See the comment above on what each of these mean. - enum class AllocationType { - kTemporary, - kLoopCarriedDependence, - kPinned, - kPrefetch, - kUnsupported - }; - - // ToString methods for logging/debugging. - static std::string AllocationTypeToString(AllocationType allocation_type); - std::string ToString() const; - - // Returns true if memory-bound loop optimizer supports allocating this type - // of a loop value. - bool IsAllocationTypeSupported() const; - - // Sets the data members `chunks`, `alternate_memory_begin_idx_in_loop`, and - // `alternate_memory_end_idx_in_loop`. - void SetChunkPairAndInterval(EvenOddChunkPair chunk_pair, - int64_t begin_idx_in_loop, - int64_t end_idx_in_loop); - - // The HloValues that correspond to this LoopValue. - std::vector hlo_values; - // The position in the header, if any. - std::optional header_position; - // The loop index and position in the previous and current iterations. - std::vector> prev_iteration_positions; - std::vector> loop_positions; - // The loop index and use in the current and next iterations. - std::vector> loop_uses; - std::vector> next_iteration_uses; - // The allocation type. - AllocationType allocation_type; - // Size of this tensor. - int64_t size; - // The default memory bandwidth savings were we to successfully put this in - // the alternate memory using the allocation type, in bytes. - float savings; - // The savings divided by the size. This is typically 2 for temporary - // allocations (skip a write and a read to the default memory). More complex - // production/consumption patterns may result in higher or lower values. We - // use this value to sort LoopValues so that the algorithm can prioritize - // allocating the buffers with the highest savings per byte to the alternate - // memory. - float savings_per_byte; - // The optimized AllocationSequence. - AllocationSequence allocations; - // Chunks for even and odd iterations. If a loop value is double buffered - // then it must have different chunks for even and odd iterations. - EvenOddChunkPair chunks; - // Begin index of loop value in alternate memory. - // REQUIRES: - // * (-loop_size) <= alternate_memory_begin_idx_in_loop - // * alternate_memory_begin_idx_in_loop < loop_size - std::optional alternate_memory_begin_idx_in_loop = std::nullopt; - // End index of loop value in alternate memory. - // REQUIRES: - // * 0 <= alternate_memory_end_idx_in_loop - // * alternate_memory_end_idx_in_loop < 2*loop_size - std::optional alternate_memory_end_idx_in_loop = std::nullopt; - }; - - // Factory method to create and initialize a MemoryBoundLoopOptimizer. - static absl::StatusOr> Create( - int loop_start, int loop_end, const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis, const Options& options); - - // Optimize the loop. Initialize must be called first. - void Optimize(); - - // Calculate the steady-state execution time of one loop iteration using the - // allocation decisions so far. - float CalculateExecutionTime() const; - - // Return the LoopValues. - const std::vector& loop_values() const { return loop_values_; } - std::vector& loop_values() { return loop_values_; } - - // Return the remaining memory vector for each point in time in the loop using - // the allocation decisions so far. - std::vector RemainingMemory() const { - return heap_.RemainingMemoryByTime(); - } - - int64_t MaxAlternateMemoryUsed() const { - return heap_.LastMemoryOffsetOccupied(); - } - - std::string MemoryUsageToAsciiArt() const { - return heap_.MemoryUsageToAsciiArt(); - } - - // The loop start, end, and size accessors. - int loop_start() const { return loop_start_; } - int loop_end() const { return loop_end_; } - int loop_size() const { return loop_size_; } - - private: - // Temporary data structures used by the AllocatePrefetch function. - struct AllocatePrefetchesContext { - // The values that are requested to be prefetched. - absl::Span values; - - // A list of indices into values array, sorted by the (descending) start - // time of the first use. - std::vector value_indices; - - // Default memory remaining bandwidths assuming all prefetches succeeded. - std::vector bandwidth_idle_times; - }; - - MemoryBoundLoopOptimizer( - int loop_start, int loop_end, uint64_t alternate_memory_size, - const MemoryBoundLoopOptimizerOptions& options, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis_, - const CostAnalysis& cost_analysis, - const BufferValue::SizeFunction* absl_nonnull size_function, - const ReservedScopedMemoryFunction& reserved_scoped_memory_fn, - int64_t alignment_in_bytes); - - // Initializes the data structures used by the optimizer. - absl::Status Initialize(); - - // Given an HloBuffer object, determines if this buffer represents a LoopValue - // that can be optimized by the optimizer, and if so it adds a LoopValue to - // the back of loop_values_ that represents the HloBuffer. Otherwise, no new - // LoopValue is added to loop_values_. - void MaybeCreateLoopValue(const HloBuffer& buffer, - const HloComputation* loop_computation); - - // Sort LoopValues by savings_per_byte. - void SortLoopValues(); - - // After allocation finishes, we fix up by creating Allocation objects to any - // LoopValues that didn't get alternate memory allocations. - void PostProcess(); - - // Allocate LoopValues by dispatching to the correct Allocate method. - void AllocateLoopValues(); - - // Perform allocation type kTemporary. Return true if successful. - bool AllocateTemporary(LoopValue& value); - - // Perform allocation type kPinned. Return true if successful. - bool AllocatePinned(LoopValue& value); - - // Perform allocation type kPrefetch. Unlike the other Allocate methods, this - // performs allocation of multiple LoopValues in order to consider the effect - // of remaining bandwidth assuming the other prefetches were successful. - // Return true if successful. - bool AllocatePrefetches(absl::Span values); - - // Allocate one prefetch for the loop value index that corresponds to - // context.context.values. Returns true if successful. - bool AllocatePrefetch(int value_index, AllocatePrefetchesContext& context); - - // Keeps track of successful allocation of all uses and positions of this - // LoopValue. - void AddAllLoopPositionsAndUses(LoopValue& value, - bool allocate_next_iteration_uses); - - // Returns the default memory bandwidth idle time at the index. - float GetBandwidthIdleTime(int idx) const; - - // Returns the default memory bandwidth idle time at the index assuming the - // given uses and positions got alternate memory allocations. - float GetBandwidthIdleTime( - int idx, - const absl::flat_hash_map>>& - additional_uses_in_alternate_mem, - const absl::flat_hash_map>& - additional_positions_in_alternate_mem) const; - - // Returns the instruction elapsed at the index. - float GetInstructionElapsed(int idx) const; - - int loop_start_; - int loop_end_; - int loop_size_; - uint64_t alternate_memory_size_; - MemoryBoundLoopOptimizerOptions options_; - const HloLiveRange& hlo_live_range_; - const HloAliasAnalysis& alias_analysis_; - const CostAnalysis& cost_analysis_; - const BufferValue::SizeFunction* absl_nonnull size_function_; - - absl::flat_hash_map instructions_in_loop_; - absl::flat_hash_map - instructions_in_prev_iteration_; - absl::flat_hash_map - instructions_in_next_iteration_; - std::vector loop_values_; - absl::flat_hash_map>> - uses_in_alternate_mem_; - absl::flat_hash_map> - positions_in_alternate_mem_; - const ReservedScopedMemoryFunction& reserved_scoped_memory_fn_; - - // The heap used to allocate loop values. Since some loop values can be double - // buffered, between successive iterations, they must have different chunks - // for even and odd iterations. We model 4 iterations of the loop to allocate - // the loop values to alternate memory so we can model the buffers that cross - // one or two loop boundaries. The allocations in the 2nd and 3rd iterations - // represent the actual memory view. The 0th and 1st iteration serve to - // account for allocations, whose buffers cross one or two loop boundaries, - // into the 2nd and 3rd iterations. - LoopOptimizerBestFitHeap heap_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_MEMORY_BOUND_LOOP_OPTIMIZER_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc deleted file mode 100644 index c3a9b452aeab2a..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc +++ /dev/null @@ -1,1746 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/memory_bound_loop_optimizer.h" - -#include -#include -#include -#include -#include -#include - -#include -#include -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/strings/match.h" -#include "absl/strings/numbers.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_replace.h" -#include "absl/strings/string_view.h" -#include "re2/re2.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/hlo/testlib/verified_hlo_module.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/buffer_value.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/buffer_interval_comparator.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/service/memory_space_assignment/prefetch_interval_picker.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/status_macros.h" -#include "xla/util.h" -#include "xla/xla_data.pb.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -using ::testing::ContainerEq; -using ::testing::HasSubstr; - -int64_t SizeFunction(const BufferValue& value) { - return HloCostAnalysis::DefaultShapeSize(value.shape()); -} - -int64_t ReservedScopedMemoryFn( - const HloInstruction* instruction, - const absl::flat_hash_set>& - operands_in_alternate_memory, - const absl::flat_hash_set& outputs_in_alternate_memory) { - return 0; -} - -class LoopOptimizerBestFitHeapTest : public ::testing::Test { - public: - LoopOptimizerBestFitHeapTest() - : heap_(/*size_limit_per_heap=*/64, /*loop_size=*/6, - /*alignment_in_bytes=*/8) {} - - bool IsAllocateSameEvenAndOddBetweenSuccessful(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size) { - EvenOddChunkPair chunks = heap_.AllocateSameEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, size); - return chunks.HasValues(); - } - - bool CanFindSameEvenAndOddAllocationBetween(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size) { - EvenOddChunkPair chunks = heap_.FindSameEvenAndOddAllocationBetween( - begin_idx_in_loop, end_idx_in_loop, size); - return chunks.HasValues(); - } - - bool IsAllocateEvenAndOddBetweenSuccessful(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size) { - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween( - begin_idx_in_loop, end_idx_in_loop, size); - return chunks.HasValues(); - } - - bool CanFindEvenAndOddAllocationBetween(int64_t begin_idx_in_loop, - int64_t end_idx_in_loop, - int64_t size) { - EvenOddChunkPair chunks = heap_.FindEvenAndOddAllocationBetween( - begin_idx_in_loop, end_idx_in_loop, size); - return chunks.HasValues(); - } - - std::string GetMemoryUsageAsciiArt() { return heap_.MemoryUsageToAsciiArt(); } - - protected: - LoopOptimizerBestFitHeap heap_; -}; - -TEST_F(LoopOptimizerBestFitHeapTest, TestAllocateSameEvenAndOddBetween) { - EXPECT_TRUE(IsAllocateSameEvenAndOddBetweenSuccessful(3, 8, 16)); - EXPECT_TRUE(IsAllocateSameEvenAndOddBetweenSuccessful(-3, 2, 16)); - EXPECT_TRUE(IsAllocateSameEvenAndOddBetweenSuccessful(0, 2, 16)); - EXPECT_TRUE(IsAllocateSameEvenAndOddBetweenSuccessful(3, 5, 16)); - EXPECT_EQ(heap_.LastMemoryOffsetOccupied(), 48); - EXPECT_TRUE(IsAllocateSameEvenAndOddBetweenSuccessful(0, 5, 16)); - EXPECT_FALSE(IsAllocateSameEvenAndOddBetweenSuccessful(0, 5, 16)); - EXPECT_EQ(heap_.LastMemoryOffsetOccupied(), 64); - EXPECT_THAT(heap_.RemainingMemoryByTime(), - ContainerEq(std::vector{0, 0, 0, 0, 0, 0})); - std::string memory_usage = heap_.MemoryUsageToAsciiArt(2, 3); - // Expected memory usage ascii art string - - // Memory map for time: [12,23], memory_block_size: 16, group_size: 6 - // - // ###### ###### 64 - // ###### ###### 48 - // ###### ###### 32 - // ###### ###### 16 - // 234567 890123 - EXPECT_THAT(memory_usage, HasSubstr("Memory map for time: [12,23], " - "memory_block_size: 16, group_size: 6")); - EXPECT_THAT(memory_usage, HasSubstr("###### ###### 64")); - EXPECT_THAT(memory_usage, HasSubstr("###### ###### 48")); - EXPECT_THAT(memory_usage, HasSubstr("###### ###### 32")); - EXPECT_THAT(memory_usage, HasSubstr("###### ###### 16")); - EXPECT_THAT(memory_usage, HasSubstr("234567 890123")); -} - -TEST_F(LoopOptimizerBestFitHeapTest, TestAllocateEvenAndOddBetween) { - EXPECT_TRUE(IsAllocateEvenAndOddBetweenSuccessful(3, 11, 16)); - EXPECT_EQ(heap_.LastMemoryOffsetOccupied(), 32); - EXPECT_TRUE(IsAllocateEvenAndOddBetweenSuccessful(-3, 8, 16)); - EXPECT_EQ(heap_.LastMemoryOffsetOccupied(), 64); - EXPECT_THAT(heap_.RemainingMemoryByTime(), - ContainerEq(std::vector{16, 16, 16, 0, 0, 0})); - std::string memory_usage = heap_.MemoryUsageToAsciiArt(); - // Expected memory usage ascii art string - - // Memory map for time: [0,35], memory_block_size: 16, group_size: 6 - // - // ...... ...### ###### ###### ###### ###... 64 - // ...### ###### ###### ###### ###... ...... 48 - // ...... ...### ###### ...### ###### ...... 32 - // ...### ###### ...### ###### ...... ...... 16 - // 012345 678901 234567 890123 456789 012345 - EXPECT_THAT( - memory_usage, - HasSubstr( - "Memory map for time: [0,35], memory_block_size: 16, group_size: 6")); - EXPECT_THAT(memory_usage, - HasSubstr("...... ...### ###### ###### ###### ###... 64")); - EXPECT_THAT(memory_usage, - HasSubstr("...### ###### ###### ###### ###... ...... 48")); - EXPECT_THAT(memory_usage, - HasSubstr("...... ...### ###### ...### ###### ...... 32")); - EXPECT_THAT(memory_usage, - HasSubstr("...### ###### ...### ###### ...... ...... 16")); - EXPECT_THAT(memory_usage, - HasSubstr("012345 678901 234567 890123 456789 012345")); -} - -TEST_F(LoopOptimizerBestFitHeapTest, TestRemoveChunk) { - EvenOddChunkPair chunks = heap_.AllocateEvenAndOddBetween(3, 11, 16); - EXPECT_TRUE(chunks.HasValues()); - EvenOddChunkPair second_chunks = heap_.AllocateEvenAndOddBetween(-3, 8, 16); - EXPECT_TRUE(second_chunks.HasValues()); - EXPECT_THAT(heap_.RemainingMemoryByTime(), - ContainerEq(std::vector{16, 16, 16, 0, 0, 0})); - EXPECT_EQ(heap_.LastMemoryOffsetOccupied(), 64); - std::string memory_usage = heap_.MemoryUsageToAsciiArt(2, 3); - // Expected memory usage ascii art string - - // Memory map for time: [12,23], memory_block_size: 16, group_size: 6 - // - // ###### ###### 64 - // ###### ###### 48 - // ###### ...### 32 - // ...### ###### 16 - // 234567 890123 - EXPECT_THAT(memory_usage, HasSubstr("Memory map for time: [12,23], " - "memory_block_size: 16, group_size: 6")); - EXPECT_THAT(memory_usage, HasSubstr("###### ###### 64")); - EXPECT_THAT(memory_usage, HasSubstr("###### ###### 48")); - EXPECT_THAT(memory_usage, HasSubstr("###### ...### 32")); - EXPECT_THAT(memory_usage, HasSubstr("...### ###### 16")); - EXPECT_THAT(memory_usage, HasSubstr("234567 890123")); - // We must 16 bytes of free memory in [0,2] with different offsets in even and - // odd iterations. - EXPECT_TRUE(CanFindEvenAndOddAllocationBetween(0, 2, 16)); - // We must not find 16 bytes of free memory in [0,2] with same offsets in even - // and odd iterations. - EXPECT_FALSE(IsAllocateSameEvenAndOddBetweenSuccessful(0, 2, 16)); - EXPECT_FALSE(CanFindEvenAndOddAllocationBetween(0, 11, 16)); - heap_.RemoveEvenOddChunkPair(3, 11, chunks); - // We must find 16 bytes of free memory spanning 2 loop iterations with - // different offsets in even and odd iterations. It does not matter what time - // is picked as the start time as long as the span is less than or equal to 2 - // iterations. - EXPECT_TRUE(CanFindEvenAndOddAllocationBetween(0, 11, 16)); - EXPECT_TRUE(CanFindEvenAndOddAllocationBetween(-3, 8, 16)); - // We must find 32 bytes of free memory less than or equal to one iteration - // with different offsets in even and odd iterations. It does not matter what - // time is picked as the start time. - EXPECT_TRUE(CanFindEvenAndOddAllocationBetween(0, 5, 32)); - EXPECT_TRUE(CanFindEvenAndOddAllocationBetween(-1, 4, 32)); - EXPECT_TRUE(CanFindEvenAndOddAllocationBetween(2, 7, 32)); - // Spans more than one iteration. - EXPECT_FALSE(CanFindEvenAndOddAllocationBetween(0, 6, 32)); - // We must be able to find 32 bytes of free memory spanning less than or equal - // to one iteration with same offsets in even and odd iterations. It does not - // matter what time is picked as the start time. - EXPECT_TRUE(CanFindSameEvenAndOddAllocationBetween(0, 5, 32)); - EXPECT_TRUE(CanFindSameEvenAndOddAllocationBetween(-1, 4, 32)); - EXPECT_TRUE(CanFindSameEvenAndOddAllocationBetween(2, 7, 32)); - std::string updated_memory_usage = heap_.MemoryUsageToAsciiArt(2, 3); - // Expected updated memory usage ascii art string - - // Memory map for time: [12,23], memory_block_size: 16, group_size: 6 - // - // ###### ###### 64 - // ###### ###### 48 - // ...... ...... 32 - // ...... ...... 16 - // 234567 890123 - EXPECT_THAT(updated_memory_usage, - HasSubstr("Memory map for time: [12,23], " - "memory_block_size: 16, group_size: 6")); - EXPECT_THAT(updated_memory_usage, HasSubstr("###### ###### 64")); - EXPECT_THAT(updated_memory_usage, HasSubstr("###### ###### 48")); - EXPECT_THAT(updated_memory_usage, HasSubstr("...... ...... 32")); - EXPECT_THAT(updated_memory_usage, HasSubstr("...... ...... 16")); - EXPECT_THAT(updated_memory_usage, HasSubstr("234567 890123")); - heap_.RemoveEvenOddChunkPair(-3, 8, second_chunks); - EXPECT_EQ(heap_.LastMemoryOffsetOccupied(), 0); -} - -class MemoryBoundLoopOptimizerTest : public HloHardwareIndependentTestBase { - public: - MemoryBoundLoopOptimizerTest() = default; - - protected: - const int64_t kAlternateMemorySpace = 1; - const int64_t kDefaultMemorySpace = 0; - - absl::Status Initialize(const HloModule* module, - uint64_t alternate_memory_size = 256) { - HloCostAnalysis::Options options; - MemoryBoundLoopOptimizerOptions optimizer_options; - optimizer_options.set_enabled(true); - optimizer_options.set_desired_copy_ratio(0.7); - optimizer_options.set_allow_unsatisfied_fully_pipelined_prefetch(false); - optimizer_options.set_min_num_iterations(3.0); - options_.memory_bound_loop_optimizer_options = optimizer_options; - cost_analysis_options_.alternate_mem_read_bandwidth_bytes_per_second = 128; - cost_analysis_options_.alternate_mem_write_bandwidth_bytes_per_second = 128; - cost_analysis_options_.default_mem_bandwidth_bytes_per_second = 32; - cost_analysis_options_.pipeline_overhead_window_size_mib = 1; - options.set_flops_per_second(16); - options.set_bytes_per_second(32); - options.set_transcendentals_per_second(16); - hlo_cost_analysis_ = std::make_unique(options); - hlo_cost_analysis_wrapper_ = - std::make_unique(*hlo_cost_analysis_); - op_cost_manager_ = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper_), - /*enable_cache=*/false)); - ABSL_ASSIGN_OR_RETURN(alias_analysis_, - HloAliasAnalysis::Run(module, &alias_info_)); - ABSL_ASSIGN_OR_RETURN( - cost_analysis_, - CostAnalysis::Create(*op_cost_manager_, cost_analysis_options_, - &alias_info_, *module, alias_analysis_.get())); - ABSL_ASSIGN_OR_RETURN(live_range_, - HloLiveRange::Run(module->schedule(), *alias_analysis_, - module->entry_computation())); - return absl::OkStatus(); - } - - absl::StatusOr CreateOptimizer( - int loop_start, int loop_end, const HloModule* module, - uint64_t alternate_memory_size = 256, - const ReservedScopedMemoryFunction& reserved_scoped_memory_fn = - ReservedScopedMemoryFn) { - ABSL_RETURN_IF_ERROR(Initialize(module, alternate_memory_size)); - MemoryBoundLoopOptimizerOptions optimizer_options; - optimizer_options.set_enabled(true); - optimizer_options.set_desired_copy_ratio(0.7); - optimizer_options.set_allow_unsatisfied_fully_pipelined_prefetch(false); - Options options; - options.max_size_in_bytes = alternate_memory_size; - options.alignment_in_bytes = 8; - options.alternate_memory_space = kAlternateMemorySpace; - options.cost_analysis = cost_analysis_.get(); - options.size_fn = SizeFunction; - options.reserved_scoped_memory_fn = reserved_scoped_memory_fn; - options.memory_bound_loop_optimizer_options = optimizer_options; - ABSL_ASSIGN_OR_RETURN(optimizer_, MemoryBoundLoopOptimizer::Create( - loop_start, loop_end, *live_range_, - *alias_analysis_, options)); - return optimizer_.get(); - } - - absl::StatusOr> ParseAndCreateOptimizer( - absl::string_view hlo_loop_str, uint64_t alternate_memory_size, - int& loop_start_idx, MemoryBoundLoopOptimizer** optimizer, - const ReservedScopedMemoryFunction& reserved_scoped_memory_fn = - ReservedScopedMemoryFn) { - int loop_end_idx; - ABSL_ASSIGN_OR_RETURN( - std::string module_str, - ParseAndCreateModuleString(hlo_loop_str, loop_start_idx, loop_end_idx)); - ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, - ParseAndReturnVerifiedModule(module_str)); - ABSL_ASSIGN_OR_RETURN( - *optimizer, - CreateOptimizer(loop_start_idx, loop_end_idx, module.get(), - alternate_memory_size, reserved_scoped_memory_fn)); - return std::move(module); - } - - // Parse a loop string description like the following: - // $op0 = f32[1,4] add(f32[1,4] $param0, f32[1,4] $prev_op4) - // $op1 = f32[8,4] add(f32[8,4] $param1, f32[8,4] $prev_op3) - // $op2 = f32[1,4] add(f32[1,4] $param2, f32[1,4] $op0) - // $op3 = f32[8,4] add(f32[8,4] $param3, f32[8,4] $op1) - // $op4 = f32[1,4] add(f32[1,4] $param4, f32[1,4] $op2) - absl::StatusOr ParseAndCreateModuleString( - absl::string_view hlo_loop_str, int& loop_start_idx, int& loop_end_idx) { - // Parse op name and types first. - RE2 op_re("\\$op([0-9]+) += +(\\S+).*"); - std::vector ops; - std::vector op_types; - int begin_pos = 0; - absl::string_view submatch[3]; - while (op_re.Match(hlo_loop_str, begin_pos, hlo_loop_str.size(), - RE2::UNANCHORED, submatch, /*nsubmatch=*/3)) { - for (int i = 0; i < 3; ++i) { - if (submatch[i].data() == nullptr) { - VLOG(4) << "Submatch[" << i << "] = nullptr"; - } else { - VLOG(4) << "Submatch[" << i << "] = " << submatch[i] - << " (idx: " << (submatch[i].data() - hlo_loop_str.data()) - << ")"; - } - } - int op_num; - if (!absl::SimpleAtoi(submatch[1], &op_num)) { - return InvalidArgument("Op name expects to contain a number, found %s.", - submatch[1]); - } - if (op_num != ops.size()) { - return InvalidArgument("Op number expected to be %d found %d.", - op_types.size(), op_num); - } - ops.push_back(submatch[0]); - op_types.push_back(submatch[2]); - begin_pos = submatch[0].data() - hlo_loop_str.data() + submatch[0].size(); - } - - RE2 param_re("([[:alnum:]]+\\[\\S*\\]) +\\$param([0-9]+)"); - std::vector param_types; - begin_pos = 0; - while (param_re.Match(hlo_loop_str, begin_pos, hlo_loop_str.size(), - RE2::UNANCHORED, submatch, /*nsubmatch=*/3)) { - for (int i = 0; i < 3; ++i) { - if (submatch[i].data() == nullptr) { - VLOG(4) << "Submatch[" << i << "] = nullptr"; - } else { - VLOG(4) << "Submatch[" << i << "] = " << submatch[i] - << " (idx: " << (submatch[i].data() - hlo_loop_str.data()) - << ")"; - } - } - int param_num; - if (!absl::SimpleAtoi(submatch[2], ¶m_num)) { - return InvalidArgument( - "Param name expects to contain a number, found %s.", submatch[2]); - } - while (param_num >= param_types.size()) { - param_types.push_back({}); - } - param_types[param_num] = submatch[1]; - - begin_pos = submatch[0].data() - hlo_loop_str.data() + submatch[0].size(); - } - - RE2 root_re("ROOT \\$root += +tuple\\((.*)\\)"); - absl::string_view root_values; - if (root_re.Match(hlo_loop_str, 0, hlo_loop_str.size(), RE2::UNANCHORED, - submatch, /*nsubmatch=*/2)) { - for (int i = 0; i < 2; ++i) { - if (submatch[i].data() == nullptr) { - VLOG(4) << "Submatch[" << i << "] = nullptr"; - } else { - VLOG(4) << "Submatch[" << i << "] = " << submatch[i] - << " (idx: " << (submatch[i].data() - hlo_loop_str.data()) - << ")"; - } - } - root_values = submatch[1]; - } - - for (absl::string_view op_type : op_types) { - VLOG(4) << "op_type: " << op_type; - } - for (absl::string_view param_type : param_types) { - VLOG(4) << "param_type: " << param_type; - } - - std::string hlo_string = R"( -HloModule module, is_scheduled=true - -ENTRY Entry { -)"; - int total_instructions = 0; - for (absl::string_view param_prefix : {"prev_", "", "next_"}) { - for (int i = 0; i < param_types.size(); ++i) { - int parameter_number = total_instructions; - absl::StrAppend(&hlo_string, " ", param_prefix, "param", i, " = ", - param_types[i], " parameter(", parameter_number, - ") // ", total_instructions++, "\n"); - } - } - - for (int i = 0; i < op_types.size(); ++i) { - int parameter_number = total_instructions; - absl::StrAppend(&hlo_string, " ", "prev_prev_op", i, " = ", op_types[i], - " parameter(", parameter_number, ") // ", - total_instructions++, "\n"); - } - - std::string new_root_values; - auto print_ops = - [&](const std::vector>& - replacements) { - for (int i = 0; i < ops.size(); ++i) { - absl::StrAppend(&hlo_string, " ", - absl::StrReplaceAll(ops[i], replacements), " // ", - total_instructions++, "\n"); - } - if (!root_values.empty()) { - absl::StrAppend(&new_root_values, - new_root_values.empty() ? "" : ", ", - absl::StrReplaceAll(root_values, replacements)); - } - }; - - std::vector> - prev_replacements; - prev_replacements.push_back({"$prev_op", "prev_prev_op"}); - prev_replacements.push_back({"$op", "prev_op"}); - prev_replacements.push_back({"$param", "prev_param"}); - absl::StrAppend(&hlo_string, " // Prev iteration body:\n"); - print_ops(prev_replacements); - - loop_start_idx = total_instructions; - std::vector> replacements; - replacements.push_back({"$", ""}); - absl::StrAppend(&hlo_string, " // Loop body:\n"); - print_ops(replacements); - loop_end_idx = total_instructions; - - std::vector> - next_replacements; - next_replacements.push_back({"$prev_op", "op"}); - next_replacements.push_back({"$op", "next_op"}); - next_replacements.push_back({"$param", "next_param"}); - absl::StrAppend(&hlo_string, " // Next iteration body:\n"); - print_ops(next_replacements); - - absl::StrAppend(&hlo_string, " ROOT root = tuple(", new_root_values, - ")\n"); - absl::StrAppend(&hlo_string, "}"); - - VLOG(1) << hlo_string; - return hlo_string; - } - - absl::StatusOr> RunMsa( - HloModule* module, uint64_t alternate_memory_size = 256) { - options_.max_size_in_bytes = alternate_memory_size; - options_.alignment_in_bytes = 8; - options_.verify = true; - - options_.alternate_memory_space = kAlternateMemorySpace; - - if (!cost_analysis_) { - ABSL_RETURN_IF_ERROR(Initialize(module, alternate_memory_size)); - } - CostAnalysis::Cache cache; - MemoryBoundednessBufferIntervalComparator comparator(*cost_analysis_, - &cache); - options_.buffer_interval_comparator = &comparator; - CostAnalysisPrefetchIntervalPicker prefetch_interval_picker( - CostAnalysisPrefetchIntervalPicker( - *cost_analysis_, /*min_overlap_to_async_copy_ratio=*/0.8, - /*preferred_overlap_to_async_copy_ratio=*/1.5, - /*max_overlap_to_mem_size_async_copy_ratio=*/10.0, - /*mem_size_bytes=*/alternate_memory_size)); - options_.prefetch_interval_picker = &prefetch_interval_picker; - - auto size_fn = [](const BufferValue& buffer) { - return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8); - }; - options_.size_fn = size_fn; - - auto is_allowed_in_alternate_mem = [](const HloValue& value) { - // Check if the value belongs to the entry computation. - HloInstruction* instruction = value.instruction(); - HloComputation* computation = instruction->parent(); - bool in_entry_computation = - (computation == computation->parent()->entry_computation()); - if (in_entry_computation && - instruction->opcode() == HloOpcode::kParameter) { - return false; - } - return true; - }; - options_.is_allowed_in_alternate_mem_fn = is_allowed_in_alternate_mem; - options_.max_outstanding_prefetches = -1; - options_.max_outstanding_evictions = -1; - options_.cost_analysis = cost_analysis_.get(); - - std::unique_ptr preset_assignments = - MemorySpaceAssignment::Run(module, *live_range_, *alias_analysis_, - &alias_info_, options_) - .value(); - return preset_assignments; - } - - absl::Status VerifyMsaEquivalence( - HloModule* module, bool expect_unsupported_allocations = false) { - // Create a map indexed by instruction number and operand number. - absl::flat_hash_map, const Allocation*> allocation_map; - for (const MemoryBoundLoopOptimizer::LoopValue& value : - optimizer_->loop_values()) { - // Skip verification for unsupported allocations as they will go through - // the usual MSA algorithm and may actually get an alternate memory - // allocation. - if (!value.IsAllocationTypeSupported()) { - continue; - } - for (const auto& allocation : value.allocations) { - for (const HloUse& use : allocation->uses()) { - absl::string_view inst_name = use.instruction->name(); - TF_RET_CHECK(absl::StartsWith(inst_name, "op")); - int inst_number; - TF_RET_CHECK(absl::SimpleAtoi(inst_name.substr(2), &inst_number)); - allocation_map[{inst_number, use.operand_number}] = allocation.get(); - } - } - } - - auto get_inst_prefix_in_iter = [](int iteration) { - switch (iteration) { - case 0: - return "prev_"; - case 1: - return ""; - case 2: - return "next_"; - default: - LOG(FATAL) << "Invalid iteration " << iteration; - return "INVALID"; - } - }; - - ABSL_ASSIGN_OR_RETURN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module, &alias_info_)); - ABSL_ASSIGN_OR_RETURN(std::unique_ptr live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - const auto& flattened_instructions = - live_range->flattened_instruction_sequence().instructions(); - for (int iteration = 1; iteration < 3; ++iteration) { - for (int inst_number = 0; inst_number < optimizer_->loop_size(); - ++inst_number) { - HloInstruction* inst = FindInstruction( - module, absl::StrCat(get_inst_prefix_in_iter(iteration), "op", - inst_number)); - for (int operand_number = 0; operand_number < 2; ++operand_number) { - const HloInstruction* operand = inst->operand(operand_number); - LOG(INFO) << inst->name() << ", operand " << operand_number; - if (!allocation_map.contains({inst_number, operand_number})) { - TF_RET_CHECK(expect_unsupported_allocations); - continue; - } - const Allocation* allocation = - allocation_map.at({inst_number, operand_number}); - if (!allocation->is_copy_allocation()) { - // We don't expect a prefetch here. - EXPECT_NE(operand->opcode(), HloOpcode::kCopyDone); - int expected_memory_space = - allocation->memory_space() == MemorySpace::kDefault - ? kDefaultMemorySpace - : kAlternateMemorySpace; - EXPECT_EQ(operand->shape().layout().memory_space(), - expected_memory_space); - } else { - EXPECT_EQ(allocation->memory_space(), MemorySpace::kAlternate); - TF_RET_CHECK(operand->opcode() == HloOpcode::kCopyDone); - const CopyAllocation* copy_allocation = - static_cast(allocation); - if (copy_allocation->copy_done_schedule_before() != inst_number) { - // The only case where the copy done schedule before is not the - // same as this use would be that this use is not the first use of - // the copy allocation. - EXPECT_NE(allocation->uses().front(), - (HloUse{inst, operand_number})); - continue; - } - int expected_copy_start_iteration = iteration; - if (copy_allocation->copy_start_schedule_after() == - optimizer_->loop_size() && - copy_allocation->copy_done_schedule_before() == 0) { - expected_copy_start_iteration -= 2; - } else if (copy_allocation->copy_start_schedule_after() + 1 >= - copy_allocation->copy_done_schedule_before()) { - expected_copy_start_iteration -= 1; - } - - if (expected_copy_start_iteration >= 0) { - const HloInstruction* expected_copy_start_schedule_after = - FindInstruction( - module, - absl::StrCat( - get_inst_prefix_in_iter( - expected_copy_start_iteration), - "op", copy_allocation->copy_start_schedule_after())); - LOG(INFO) << "Expected copy start schedule after: " - << expected_copy_start_schedule_after->name(); - const HloInstruction* copy_start = operand->operand(0); - TF_RET_CHECK(copy_start->opcode() == HloOpcode::kCopyStart); - // Find the instruction before this copy start that is not an - // async copy or gte or parameter. - int copy_start_idx = - live_range->instruction_schedule().at(copy_start); - const HloInstruction* copy_start_schedule_after = nullptr; - for (int i = copy_start_idx - 1; i >= 0; --i) { - HloOpcode opcode = flattened_instructions.at(i)->opcode(); - if (opcode != HloOpcode::kCopyStart && - opcode != HloOpcode::kCopyDone && - opcode != HloOpcode::kGetTupleElement && - opcode != HloOpcode::kParameter) { - copy_start_schedule_after = flattened_instructions.at(i); - break; - } - } - TF_RET_CHECK(copy_start_schedule_after != nullptr); - EXPECT_EQ(copy_start_schedule_after, - expected_copy_start_schedule_after); - } - } - } - } - } - return absl::OkStatus(); - } - - AliasInfo alias_info_; - - private: - Options options_; - CostAnalysisOptions cost_analysis_options_; - std::unique_ptr hlo_cost_analysis_; - std::unique_ptr hlo_cost_analysis_wrapper_; - std::unique_ptr op_cost_manager_; - std::unique_ptr cost_analysis_; - std::unique_ptr alias_analysis_; - std::unique_ptr live_range_; - std::unique_ptr optimizer_; -}; - -TEST_F(MemoryBoundLoopOptimizerTest, SimplePrefetch) { - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op3, f32[1,4] $prev_op4) - $op1 = f32[1,4] add(f32[1,4] $prev_op4, f32[1,4] $op0) - $op2 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op1) - $op3 = f32[1,4] add(f32[1,4] $op1, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $param0, f32[1,4] $op3) - ROOT $root = tuple($op4, $param0) - )"; - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - // Although alternate_memory_size=64 is minimum memory needed to fit the copy - // of param0 with desired copy ratio. alternate_memory_size=80 memory will - // ensure complete copy of param0 to alternate memory. - int64_t alternate_memory_size = 80; - ASSERT_OK_AND_ASSIGN( - auto module, ParseAndCreateOptimizer(hlo_loop_str, alternate_memory_size, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - absl::flat_hash_set seen_uses; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - LOG(INFO) << loop_value.ToString(); - if (loop_value.hlo_values.front() - ->defining_position() - .instruction->name() == "param0") { - EXPECT_TRUE(loop_value.allocations.back()->is_copy_allocation()); - } - for (const auto& allocation : loop_value.allocations) { - for (const HloUse& use : allocation->uses()) { - EXPECT_FALSE(seen_uses.contains(use)) << use.ToString(); - seen_uses.insert(use); - } - } - } - - // Ensure all of the uses in the loop have an associated use. - for (absl::string_view inst_name : {"op0", "op1", "op2", "op3", "op4"}) { - HloInstruction* inst = - module->entry_computation()->GetInstructionWithName(inst_name); - EXPECT_TRUE(seen_uses.contains(HloUse{inst, 0})) << inst_name; - EXPECT_TRUE(seen_uses.contains(HloUse{inst, 1})) << inst_name; - } - EXPECT_EQ(optimizer->CalculateExecutionTime(), 1.875); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -TEST_F(MemoryBoundLoopOptimizerTest, SimplePrefetch2) { - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op3, f32[1,4] $prev_op4) - $op1 = f32[1,4] add(f32[1,4] $prev_op4, f32[1,4] $op0) - $op2 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op1) - $op3 = f32[1,4] add(f32[1,4] $op1, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $param0, f32[1,4] $op3) - ROOT $root = tuple($op4, $param0) - )"; - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - // alternate_memory_size=64 is minimum memory needed to fit the copy of param0 - // with desired copy ratio. - int64_t alternate_memory_size = 64; - ASSERT_OK_AND_ASSIGN( - auto module, ParseAndCreateOptimizer(hlo_loop_str, alternate_memory_size, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - absl::flat_hash_set seen_uses; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - LOG(INFO) << loop_value.ToString(); - if (loop_value.hlo_values.front() - ->defining_position() - .instruction->name() == "param0") { - EXPECT_TRUE(loop_value.allocations.back()->is_copy_allocation()); - } - for (const auto& allocation : loop_value.allocations) { - for (const HloUse& use : allocation->uses()) { - EXPECT_FALSE(seen_uses.contains(use)) << use.ToString(); - seen_uses.insert(use); - } - } - } - - // Ensure all of the uses in the loop have an associated use. - for (absl::string_view inst_name : {"op0", "op1", "op2", "op3", "op4"}) { - HloInstruction* inst = - module->entry_computation()->GetInstructionWithName(inst_name); - EXPECT_TRUE(seen_uses.contains(HloUse{inst, 0})) << inst_name; - EXPECT_TRUE(seen_uses.contains(HloUse{inst, 1})) << inst_name; - } - // Check that execution time has increased to 2 since we will wait on copy - // done for param0. - EXPECT_EQ(optimizer->CalculateExecutionTime(), 2); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -// Specify a ReservedScopedMemoryFunction to the loop optimizer that causes each -// HLO to reserve the entire alternate memory. If the loop optimizer is -// correctly accounting for reserved scoped memory, it should not put any -// allocations in alternate memory, which we test. -TEST_F(MemoryBoundLoopOptimizerTest, ReservedScopedMemory) { - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op3, f32[1,4] $prev_op4) - $op1 = f32[1,4] add(f32[1,4] $prev_op4, f32[1,4] $op0) - $op2 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op1) - $op3 = f32[1,4] add(f32[1,4] $op1, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $param0, f32[1,4] $op3) - ROOT $root = tuple($op4, $param0) - )"; - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - ASSERT_OK_AND_ASSIGN( - auto module, - ParseAndCreateOptimizer( - hlo_loop_str, - /*alternate_memory_size=*/128, loop_start_idx, &optimizer, - [](const HloInstruction*, - const absl::flat_hash_set>&, - const absl::flat_hash_set&) { return 128; })); - - optimizer->Optimize(); - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - LOG(INFO) << "Loop value: " << loop_value.ToString(); - for (const auto& allocation : loop_value.allocations) { - ASSERT_NE(static_cast(allocation->memory_space()), - kAlternateMemorySpace); - } - } -} - -// Check that a spurious GetTupleElement instruction in a later iteration of a -// loop does not cause MSA to CHECK fail, when identifying loops. Prior to the -// change introduced with this test, IdentifyAndOptimizeMemoryBoundLoops() -// would recognize 4 iterations to the loop thinking that gte is a repeat of -// op2. Doing so triggers the CHECKs introduced by the change that added this -// test to fail. So, the point of this test is to verify that we do not check -// fail. -TEST_F(MemoryBoundLoopOptimizerTest, GetTupleElement) { - absl::string_view hlo_string = R"( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[1,4] parameter(0) - p1 = f32[1,4] parameter(1) - p2 = f32[1,4] parameter(2) - p3 = f32[1,4] parameter(3) - p4 = f32[1,4] parameter(4) - p5 = f32[1,4] parameter(5) - p6 = f32[1,4] parameter(6) - tupleparam = (f32[1,4], f32[1,4]) parameter(7) - - // Iteration 0 - op1 = tanh(p0) - op2 = tanh(p1) - op3 = tanh(op2) - op4 = add(op1, op3) - - // Iteration 1 - op5 = tanh(p2) - op6 = tanh(p3) - op7 = tanh(op6) - op8 = add(op5, op7) - - // Iteration 2 - op9 = tanh(p4) - op10 = tanh(p5) - op11 = tanh(op10) - op12 = add(op9, op11) - - // Not an iteration - op13 = tanh(p6) - gte = get-tuple-element(tupleparam), index=1 - op14 = tanh(gte) - op15 = tanh(op14) - op16 = add(op13, op15) - - ROOT root = tuple(tupleparam, op4, op8, op12, op16) - })"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - ASSERT_OK_AND_ASSIGN(auto preset_assignments, RunMsa(module.get())); -} - -TEST_F(MemoryBoundLoopOptimizerTest, NoAlternateMem) { - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op3, f32[1,4] $prev_op4) - $op1 = f32[1,4] add(f32[1,4] $prev_op4, f32[1,4] $op0) - $op2 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op1) - $op3 = f32[1,4] add(f32[1,4] $op1, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $param0, f32[1,4] $op3) - ROOT $root = tuple($op4, $param0) - )"; - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - // Set alternate memory size to zero so nothing should be in the alternate - // memory. We still expect to find an allocation for all uses. - ASSERT_OK_AND_ASSIGN(auto module, - ParseAndCreateOptimizer(hlo_loop_str, - /*alternate_memory_size=*/0, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - absl::flat_hash_set seen_uses; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - LOG(INFO) << loop_value.ToString(); - for (const auto& allocation : loop_value.allocations) { - EXPECT_EQ(allocation->memory_space(), MemorySpace::kDefault); - for (const HloUse& use : allocation->uses()) { - EXPECT_FALSE(seen_uses.contains(use)) << use.ToString(); - seen_uses.insert(use); - } - } - } - - // Ensure all of the uses in the loop have an associated use. - for (absl::string_view inst_name : {"op0", "op1", "op2", "op3", "op4"}) { - HloInstruction* inst = - module->entry_computation()->GetInstructionWithName(inst_name); - EXPECT_TRUE(seen_uses.contains(HloUse{inst, 0})) << inst_name; - EXPECT_TRUE(seen_uses.contains(HloUse{inst, 1})) << inst_name; - } -} - -TEST_F(MemoryBoundLoopOptimizerTest, PrefetchFifoOrderWithOverlap) { - // Test for enforcing FIFO order of prefetches. There are three parameters - // that will be prefetched (param0, param1, and param2). param2 is one eighth - // the size of the other parameters and is scheduled later in the loop. So, we - // expect the allocation algorithm to initially allocate param2's prefetch - // with a short live range (since copying it doesn't take very long), but then - // as we try to prefetch param0 and param1, we will wrap around into the - // previous iterations and would need to "early force" param2's prefetch to be - // scheduled earlier to enforce the FIFO order. - // - // alternate_mem_bytes_per_second = 128 - // default_mem_bytes_per_second = 32 - // flops_per_second = 16 - // f32[1,4] add: flops: 4, bytes: 48, compute elapsed: 0.25 - // - All default memory elapsed: 1.5 - // - All alternate memory elapsed: 0.375 - // f32[8,4] add: flops: 32, bytes: 384, compute elapsed: 2 - // - All default memory elapsed: 12 - // - All alternate memory elapsed: 3 - // f32[1,4] copy: bytes: 16, memory elapsed: 0.5 - // f32[8,4] copy: bytes: 128, memory elapsed: 4 - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op13, f32[1,4] $prev_op14) - $op1 = f32[8,4] add(f32[8,4] $param0, f32[8,4] $param1) - $op2 = f32[1,4] add(f32[1,4] $prev_op14, f32[1,4] $op0) - $op3 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $op2, f32[1,4] $op3) - $op5 = f32[1,4] add(f32[1,4] $op3, f32[1,4] $op4) - $op6 = f32[1,4] add(f32[1,4] $op4, f32[1,4] $op5) - $op7 = f32[1,4] add(f32[1,4] $op5, f32[1,4] $op6) - $op8 = f32[1,4] add(f32[1,4] $op6, f32[1,4] $op7) - $op9 = f32[1,4] add(f32[1,4] $op7, f32[1,4] $op8) - $op10 = f32[1,4] add(f32[1,4] $op8, f32[1,4] $op9) - $op11 = f32[1,4] add(f32[1,4] $op9, f32[1,4] $op10) - $op12 = f32[1,4] add(f32[1,4] $op10, f32[1,4] $op11) - $op13 = f32[1,4] add(f32[1,4] $op11, f32[1,4] $op12) - $op14 = f32[1,4] add(f32[1,4] $param2, f32[1,4] $op13) - )"; - - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - int64_t alternate_memory_size = 464; - ASSERT_OK_AND_ASSIGN( - auto module, ParseAndCreateOptimizer(hlo_loop_str, alternate_memory_size, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - // We expect the prefetches to be scheduled this way: - // - // - // param0 or param1: - // ===========> =====================================> - // param1 or param0: - // ===========> === - // ==============================================> - // param2: - // =====> ========================================> === - // 13 14| 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14| 0 1 - // prev | loop | next - // - // Temporaries: - // +======+ - // +=========+ - // +=========+ - // +======+ - // +======+ - // +======+ - // +======+ - // +======+ - // +======+ - // +======+ - // +======+ - // +======+ - // +======+ - // +===+ - // +======+ - // +=========+ - // 13 14| 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14| 0 1 - // prev | loop | next - std::vector prefetches; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - if (!loop_value.allocations.empty() && - loop_value.allocations.back()->is_copy_allocation()) { - prefetches.push_back(static_cast( - loop_value.allocations.back().get())); - } - } - EXPECT_EQ(prefetches.size(), 3); - bool seen_overlap = false; - bool seen_nonoverlap = false; - for (const CopyAllocation* prefetch : prefetches) { - const HloUse& use = *prefetch->uses().begin(); - if (use.instruction->name() == "op14") { - EXPECT_EQ(prefetch->copy_done_schedule_before(), 14); - EXPECT_EQ(prefetch->copy_start_schedule_after(), 0); - } else { - ASSERT_EQ(use.instruction->name(), "op1"); - EXPECT_EQ(prefetch->copy_done_schedule_before(), 1); - if (prefetch->copy_start_schedule_after() == 0) { - EXPECT_FALSE(seen_overlap); - seen_overlap = true; - } else { - EXPECT_GT(prefetch->copy_start_schedule_after(), 1); - EXPECT_FALSE(seen_nonoverlap); - seen_nonoverlap = true; - } - } - } - // We expect to fully saturate the default memory bandwidth. Total default - // memory accesses: - // param0 (128 B) + param1 (128 B) + op1 (128 B) + param2 (16 B) = 400 B - // execution time: - // 400 B / 32 B/s = 12.5 s. - EXPECT_EQ(optimizer->CalculateExecutionTime(), 12.5); - - // Check the memory used at each point of the loop. - std::vector remaining_memory = optimizer->RemainingMemory(); - // Time 0: 3 temporaries (16 B) + param0 (128 B) + param1 (128 B) - EXPECT_EQ(remaining_memory.at(0), - alternate_memory_size - (3 * 16 + 128 + 128)); - // Time 1: 2 temporaries (16 B) + 2*param0 (128 B) + param1 (128 B) - // + param2 (16 B) - EXPECT_EQ(remaining_memory.at(1), - alternate_memory_size - (2 * 16 + 2 * 128 + 128 + 16)); - // Times 2 and 3: 3 temporaries (16 B) + param0 (128 B) + param2 (16 B) - EXPECT_EQ(remaining_memory.at(2), - alternate_memory_size - (3 * 16 + 128 + 16)); - EXPECT_EQ(remaining_memory.at(3), - alternate_memory_size - (3 * 16 + 128 + 16)); - // Times 4 to 13: 3 temporaries (16 B) + param0 (128 B) + param1 (128 B) - // + param2 (16 B) - for (int i = 4; i <= 13; ++i) { - EXPECT_EQ(remaining_memory.at(i), - alternate_memory_size - (3 * 16 + 128 + 128 + 16)); - } - // Time 14: 2 temporaries (16 B) + param0 (128 B) + param1 (128 B) - // + param2 (16 B) - EXPECT_EQ(remaining_memory.at(14), - alternate_memory_size - (2 * 16 + 128 + 128 + 16)); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -TEST_F(MemoryBoundLoopOptimizerTest, PrefetchFifoOrderWithoutOverlap) { - // Same as the test above, except the size of alternate memory is less than - // 384, which is the minimum amount needed to keep the three 128-byte sized - // parameters alive (one of the parameters would need to be overlapped with - // the previous iteration, so counts 2X). In that case, we won't be able to - // fully saturate the bandwidth. - // - // alternate_mem_bytes_per_second = 128 - // default_mem_bytes_per_second = 32 - // flops_per_second = 16 - // f32[1,4] add: flops: 4, bytes: 48, compute elapsed: 0.25 - // - All default memory elapsed: 1.5 - // - All alternate memory elapsed: 0.375 - // f32[8,4] add: flops: 32, bytes: 384, compute elapsed: 2 - // - All default memory elapsed: 12 - // - All alternate memory elapsed: 3 - // f32[1,4] copy: bytes: 16, memory elapsed: 0.5 - // f32[8,4] copy: bytes: 128, memory elapsed: 4 - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op13, f32[1,4] $prev_op14) - $op1 = f32[8,4] add(f32[8,4] $param0, f32[8,4] $param1) - $op2 = f32[1,4] add(f32[1,4] $prev_op14, f32[1,4] $op0) - $op3 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $op2, f32[1,4] $op3) - $op5 = f32[1,4] add(f32[1,4] $op3, f32[1,4] $op4) - $op6 = f32[1,4] add(f32[1,4] $op4, f32[1,4] $op5) - $op7 = f32[1,4] add(f32[1,4] $op5, f32[1,4] $op6) - $op8 = f32[1,4] add(f32[1,4] $op6, f32[1,4] $op7) - $op9 = f32[1,4] add(f32[1,4] $op7, f32[1,4] $op8) - $op10 = f32[1,4] add(f32[1,4] $op8, f32[1,4] $op9) - $op11 = f32[1,4] add(f32[1,4] $op9, f32[1,4] $op10) - $op12 = f32[1,4] add(f32[1,4] $op10, f32[1,4] $op11) - $op13 = f32[1,4] add(f32[1,4] $op11, f32[1,4] $op12) - $op14 = f32[1,4] add(f32[1,4] $param2, f32[1,4] $op13) - )"; - - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - int64_t alternate_memory_size = 208; - ASSERT_OK_AND_ASSIGN( - auto module, ParseAndCreateOptimizer(hlo_loop_str, alternate_memory_size, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - // We expect the prefetches to be scheduled this way: - // - // - // param0 or param1: - // ===========> =====================================> - // param2: - // =====> ===============================> - // 13 14| 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14| 0 1 - // prev | loop | next - std::vector prefetches; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - if (!loop_value.allocations.empty() && - loop_value.allocations.back()->is_copy_allocation()) { - prefetches.push_back(static_cast( - loop_value.allocations.back().get())); - } - } - EXPECT_EQ(prefetches.size(), 2); - std::optional expected_op14_copy_start_time; - for (const CopyAllocation* prefetch : prefetches) { - const HloUse& use = *prefetch->uses().begin(); - if (use.instruction->name() == "op1") { - EXPECT_EQ(prefetch->copy_done_schedule_before(), 1); - EXPECT_GT(prefetch->copy_start_schedule_after(), 1); - expected_op14_copy_start_time = prefetch->copy_start_schedule_after(); - } - } - EXPECT_TRUE(expected_op14_copy_start_time.has_value()); - for (const CopyAllocation* prefetch : prefetches) { - const HloUse& use = *prefetch->uses().begin(); - if (use.instruction->name() == "op14") { - EXPECT_EQ(prefetch->copy_done_schedule_before(), 14); - EXPECT_EQ(prefetch->copy_start_schedule_after(), - *expected_op14_copy_start_time); - } - } - // We expect not to fully saturate the default memory bandwidth. - EXPECT_GT(optimizer->CalculateExecutionTime(), 12.5); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -TEST_F(MemoryBoundLoopOptimizerTest, PrefetchFifoOrderWithOverlap2) { - // Same as PrefetchFifoOrderWithOverlap, except the instructions are shifted - // earlier by one such that param0 and param1 are used by op0. This tests that - // we are accounting for overlaps for prefetches that span three iterations. - // - // alternate_mem_bytes_per_second = 128 - // default_mem_bytes_per_second = 32 - // flops_per_second = 16 - // f32[1,4] add: flops: 4, bytes: 48, compute elapsed: 0.25 - // - All default memory elapsed: 1.5 - // - All alternate memory elapsed: 0.375 - // f32[8,4] add: flops: 32, bytes: 384, compute elapsed: 2 - // - All default memory elapsed: 12 - // - All alternate memory elapsed: 3 - // f32[1,4] copy: bytes: 16, memory elapsed: 0.5 - // f32[8,4] copy: bytes: 128, memory elapsed: 4 - absl::string_view hlo_loop_str = R"( - $op0 = f32[8,4] add(f32[8,4] $param0, f32[8,4] $param1) - $op1 = f32[1,4] add(f32[1,4] $prev_op13, f32[1,4] $prev_op14) - $op2 = f32[1,4] add(f32[1,4] $prev_op14, f32[1,4] $op1) - $op3 = f32[1,4] add(f32[1,4] $op1, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $op2, f32[1,4] $op3) - $op5 = f32[1,4] add(f32[1,4] $op3, f32[1,4] $op4) - $op6 = f32[1,4] add(f32[1,4] $op4, f32[1,4] $op5) - $op7 = f32[1,4] add(f32[1,4] $op5, f32[1,4] $op6) - $op8 = f32[1,4] add(f32[1,4] $op6, f32[1,4] $op7) - $op9 = f32[1,4] add(f32[1,4] $op7, f32[1,4] $op8) - $op10 = f32[1,4] add(f32[1,4] $op8, f32[1,4] $op9) - $op11 = f32[1,4] add(f32[1,4] $op9, f32[1,4] $op10) - $op12 = f32[1,4] add(f32[1,4] $op10, f32[1,4] $op11) - $op13 = f32[1,4] add(f32[1,4] $param2, f32[1,4] $op12) - $op14 = f32[1,4] add(f32[1,4] $op12, f32[1,4] $op13) - )"; - - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - int64_t alternate_memory_size = 464; - ASSERT_OK_AND_ASSIGN( - auto module, ParseAndCreateOptimizer(hlo_loop_str, alternate_memory_size, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - // We expect the prefetches to be scheduled this way: - // - // - // param0 or param1: - // ========> =====================================> === - // param1 or param0: - // ========> ====== - // ==============================================> - // param2: - // ==> ========================================> ====== - // 13 14| 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14| 0 1 - // prev | loop | next - std::vector prefetches; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - if (!loop_value.allocations.empty() && - loop_value.allocations.back()->is_copy_allocation()) { - prefetches.push_back(static_cast( - loop_value.allocations.back().get())); - } - } - EXPECT_EQ(prefetches.size(), 3); - bool seen_overlap = false; - bool seen_nonoverlap = false; - for (const CopyAllocation* prefetch : prefetches) { - const HloUse& use = *prefetch->uses().begin(); - if (use.instruction->name() == "op13") { - EXPECT_EQ(prefetch->copy_done_schedule_before(), 13); - EXPECT_EQ(prefetch->copy_start_schedule_after(), 14); - } else { - ASSERT_EQ(use.instruction->name(), "op0"); - EXPECT_EQ(prefetch->copy_done_schedule_before(), 0); - if (prefetch->copy_start_schedule_after() == 14) { - EXPECT_FALSE(seen_overlap); - seen_overlap = true; - } else { - EXPECT_LT(prefetch->copy_start_schedule_after(), 14); - EXPECT_FALSE(seen_nonoverlap); - seen_nonoverlap = true; - } - } - } - // We expect to fully saturate the default memory bandwidth. Total default - // memory accesses: - // param0 (128 B) + param1 (128 B) + op1 (128 B) + param2 (16 B) = 400 B - // execution time: - // 400 B / 32 B/s = 12.5 s. - EXPECT_EQ(optimizer->CalculateExecutionTime(), 12.5); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -TEST_F(MemoryBoundLoopOptimizerTest, PrefetchFifoOrderWithoutOverlap2) { - // Same as PrefetchFifoOrderWithoutOverlap, except that, we reduce the size of - // alternate memory, such that only one of param0 and param1 can be - // prefetched. Additionally, we add many more small prefetches, such that, - // during the prefetch of param0 or param1, a valid copy start time is found - // with desired copy ratio, but not with complete copy ratio, early forcing - // param2. After finding a copy start time with the desired copy ratio, when - // trying to find a better copy start time with complete copy ratio, more - // prefetches are temporarily early forced, but restored to their original - // state later. - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op13, f32[1,4] $prev_op14) - $op1 = f32[8,4] add(f32[8,4] $param0, f32[8,4] $param1) - $op2 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $param9) - $op3 = f32[1,4] add(f32[1,4] $param7, f32[1,4] $param8) - $op4 = f32[1,4] add(f32[1,4] $param5, f32[1,4] $param6) - $op5 = f32[1,4] add(f32[1,4] $param4, f32[1,4] $op3) - $op6 = f32[1,4] add(f32[1,4] $op4, f32[1,4] $param3) - $op7 = f32[1,4] add(f32[1,4] $op5, f32[1,4] $op2) - $op8 = f32[1,4] add(f32[1,4] $op6, f32[1,4] $op7) - $op9 = f32[1,4] add(f32[1,4] $op7, f32[1,4] $op8) - $op10 = f32[1,4] add(f32[1,4] $op8, f32[1,4] $op9) - $op11 = f32[1,4] add(f32[1,4] $op9, f32[1,4] $op10) - $op12 = f32[1,4] add(f32[1,4] $op10, f32[1,4] $op11) - $op13 = f32[1,4] add(f32[1,4] $op11, f32[1,4] $op12) - $op14 = f32[1,4] add(f32[1,4] $param2, f32[1,4] $op13) - )"; - - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - int64_t alternate_memory_size = 384; - ASSERT_OK_AND_ASSIGN( - auto module, ParseAndCreateOptimizer(hlo_loop_str, alternate_memory_size, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - std::vector prefetches; - for (const MemoryBoundLoopOptimizer::LoopValue& loop_value : - optimizer->loop_values()) { - if (!loop_value.allocations.empty() && - loop_value.allocations.back()->is_copy_allocation()) { - prefetches.push_back(static_cast( - loop_value.allocations.back().get())); - } - } - EXPECT_EQ(prefetches.size(), 9); - for (const CopyAllocation* prefetch : prefetches) { - const HloUse& use = *prefetch->uses().begin(); - if (use.instruction->name() == "op14") { - EXPECT_EQ(prefetch->copy_done_schedule_before(), 14); - EXPECT_EQ(prefetch->copy_start_schedule_after(), 6); - } else if (use.instruction->name() == "op1") { - EXPECT_EQ(prefetch->copy_start_schedule_after(), 6); - EXPECT_EQ(prefetch->copy_done_schedule_before(), 1); - } - } - - EXPECT_NEAR(optimizer->CalculateExecutionTime(), 16.7083, 1e-3); - const std::vector& remaining_memory = optimizer->RemainingMemory(); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(0), 176); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(1), 208); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(2), 112); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(3), 112); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(4), 112); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(5), 96); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(6), 80); - EXPECT_EQ(alternate_memory_size - remaining_memory.at(7), 208); - for (int i = 8; i < 14; ++i) { - EXPECT_EQ(alternate_memory_size - remaining_memory.at(i), 192); - } - EXPECT_EQ(alternate_memory_size - remaining_memory.at(14), 176); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size - 128); -} - -TEST_F(MemoryBoundLoopOptimizerTest, OptimizerEndToEnd) { - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op13, f32[1,4] $prev_op14) - $op1 = f32[8,4] add(f32[8,4] $param0, f32[8,4] $param1) - $op2 = f32[1,4] add(f32[1,4] $prev_op14, f32[1,4] $op0) - $op3 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $op2, f32[1,4] $op3) - $op5 = f32[1,4] add(f32[1,4] $op3, f32[1,4] $op4) - $op6 = f32[1,4] add(f32[1,4] $op4, f32[1,4] $op5) - $op7 = f32[1,4] add(f32[1,4] $op5, f32[1,4] $op6) - $op8 = f32[1,4] add(f32[1,4] $op6, f32[1,4] $op7) - $op9 = f32[1,4] add(f32[1,4] $op7, f32[1,4] $op8) - $op10 = f32[1,4] add(f32[1,4] $op8, f32[1,4] $op9) - $op11 = f32[1,4] add(f32[1,4] $op9, f32[1,4] $op10) - $op12 = f32[1,4] add(f32[1,4] $op10, f32[1,4] $op11) - $op13 = f32[1,4] add(f32[1,4] $op11, f32[1,4] $op12) - $op14 = f32[1,4] add(f32[1,4] $param2, f32[1,4] $op13) - ROOT $root = tuple($op1, $op14) - )"; - - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - ASSERT_OK_AND_ASSIGN(auto module, - ParseAndCreateOptimizer(hlo_loop_str, - /*alternate_memory_size=*/1024, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - ASSERT_OK_AND_ASSIGN(auto preset_assignments, - RunMsa(module.get(), /*alternate_memory_size=*/1024)); - - ASSERT_OK(VerifyMsaEquivalence(module.get())); -} - -TEST_F(MemoryBoundLoopOptimizerTest, OptimizerEndToEndUnsupportedAllocation) { - // op2 is a loop-carried dependency, which is currently not supported. But the - // usual MSA algorithm should still be able to give it an alternate memory - // allocation. - absl::string_view hlo_loop_str = R"( - $op0 = f32[1,4] add(f32[1,4] $prev_op3, f32[1,4] $prev_op4) - $op1 = f32[8,4] add(f32[8,4] $param0, f32[8,4] $param1) - $op2 = f32[1,4] add(f32[1,4] $prev_op2, f32[1,4] $op0) - $op3 = f32[1,4] add(f32[1,4] $op0, f32[1,4] $op2) - $op4 = f32[1,4] add(f32[1,4] $op2, f32[1,4] $op3) - ROOT $root = tuple($op1, $op4) - )"; - - int loop_start_idx; - MemoryBoundLoopOptimizer* optimizer; - ASSERT_OK_AND_ASSIGN(auto module, - ParseAndCreateOptimizer(hlo_loop_str, - /*alternate_memory_size=*/1024, - loop_start_idx, &optimizer)); - - optimizer->Optimize(); - ASSERT_OK_AND_ASSIGN(auto preset_assignments, - RunMsa(module.get(), /*alternate_memory_size=*/1024)); - - ASSERT_OK(VerifyMsaEquivalence(module.get(), - /*expect_unsupported_allocations=*/true)); - - const HloInstruction* op2 = FindInstruction(module.get(), "op2"); - EXPECT_EQ(op2->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemoryBoundLoopOptimizerTest, TempAndPinnedAllocations) { - absl::string_view hlo_str = R"( - HloModule module, is_scheduled=true - - while_cond { - while_cond_param = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(while_cond_param), index=5 - } - - while_body { - while_body_param = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - pinned_prev_param0 = f32[1,4] get-tuple-element(while_body_param), index=0 - next_param0 = f32[1,4] get-tuple-element(while_body_param), index=1 - prev_prev_op3 = f32[1,4] get-tuple-element(while_body_param), index=2 - prev_prev_op4 = f32[1,4] get-tuple-element(while_body_param), index=3 - prev_op0 = f32[1,4] add(f32[1,4] prev_prev_op3, f32[1,4] prev_prev_op4) - prev_op1 = f32[1,4] add(f32[1,4] prev_prev_op4, f32[1,4] prev_op0) - prev_op2 = f32[1,4] add(f32[1,4] prev_op0, f32[1,4] prev_op1) - prev_op3 = f32[1,4] add(f32[1,4] prev_op1, f32[1,4] prev_op2) - prev_op4 = f32[1,4] multiply(f32[1,4] pinned_prev_param0, f32[1,4] prev_op3) - op0 = f32[1,4] add(f32[1,4] prev_op3, f32[1,4] prev_op4) - op1 = f32[1,4] add(f32[1,4] prev_op4, f32[1,4] op0) - op2 = f32[1,4] add(f32[1,4] op0, f32[1,4] op1) - op3 = f32[1,4] add(f32[1,4] op1, f32[1,4] op2) - op4 = f32[1,4] multiply(f32[1,4] pinned_prev_param0, f32[1,4] op3) - next_op0 = f32[1,4] add(f32[1,4] op3, f32[1,4] op4) - next_op1 = f32[1,4] add(f32[1,4] op4, f32[1,4] next_op0) - next_op2 = f32[1,4] add(f32[1,4] next_op0, f32[1,4] next_op1) - next_op3 = f32[1,4] add(f32[1,4] next_op1, f32[1,4] next_op2) - next_op4 = f32[1,4] multiply(f32[1,4] pinned_prev_param0, f32[1,4] next_op3) - p = pred[] get-tuple-element(while_body_param), index=5 - ROOT root = tuple(pinned_prev_param0, next_param0, prev_prev_op3, prev_prev_op4, next_op4, p) - } - - ENTRY entry { - p0 = f32[1,4] parameter(0) - p1 = f32[1,4] parameter(1) - p2 = f32[1,4] parameter(2) - p3 = f32[1,4] parameter(3) - p4 = pred[] parameter(4) - copy = f32[1,4] copy(p3) - tuple = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) tuple(p0, p1, p2, p3, copy, p4) - while = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) while(tuple), condition=while_cond, body=while_body - ROOT root = f32[1,4] get-tuple-element(while), index=4 - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str)); - int64_t alternate_memory_size = 80; - ASSERT_OK_AND_ASSIGN(auto optimizer, CreateOptimizer(19, 24, module.get(), - alternate_memory_size)); - optimizer->Optimize(); - - std::vector remaining_memory = optimizer->RemainingMemory(); - // Time 0: 3 temporaries (16 B) + 1 pinned (16 B) - EXPECT_EQ(remaining_memory.at(0), alternate_memory_size - (3 * 16 + 16)); - // Time 1: 3 temporaries (16 B) + 1 pinned (16 B) - EXPECT_EQ(remaining_memory.at(1), alternate_memory_size - (3 * 16 + 16)); - // Time 2: 3 temporaries (16 B) + 1 pinned (16 B) - EXPECT_EQ(remaining_memory.at(2), alternate_memory_size - (3 * 16 + 16)); - // Time 3: 3 temporaries (16 B) + 1 pinned (16 B) - EXPECT_EQ(remaining_memory.at(3), alternate_memory_size - (3 * 16 + 16)); - // Time 4: 2 temporaries (16 B) + 1 pinned (16 B) - EXPECT_EQ(remaining_memory.at(4), alternate_memory_size - (2 * 16 + 16)); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -TEST_F(MemoryBoundLoopOptimizerTest, NegativeSavingNotPinned) { - absl::string_view hlo_str = R"( - HloModule module, is_scheduled=true - - while_cond { - while_cond_param = (f32[28,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(while_cond_param), index=5 - } - - while_body { - while_body_param = (f32[28,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - pinned_prev_param0 = f32[28,4] get-tuple-element(while_body_param), index=0 - zero = s32[] constant(0) - next_param0 = f32[1,4] get-tuple-element(while_body_param), index=1 - prev_prev_op3 = f32[1,4] get-tuple-element(while_body_param), index=2 - prev_prev_op4 = f32[1,4] get-tuple-element(while_body_param), index=3 - prev_op0 = f32[1,4] add(f32[1,4] prev_prev_op3, f32[1,4] prev_prev_op4) - prev_op1 = f32[1,4] add(f32[1,4] prev_prev_op4, f32[1,4] prev_op0) - prev_op2 = f32[1,4] add(f32[1,4] prev_op0, f32[1,4] prev_op1) - prev_op3 = f32[1,4] add(f32[1,4] prev_op1, f32[1,4] prev_op2) - pinned_slice = f32[1,4] dynamic-slice(pinned_prev_param0, zero, zero), dynamic_slice_sizes={1,4} - prev_op4 = f32[1,4] multiply(f32[1,4] pinned_slice, f32[1,4] prev_op3) - op0 = f32[1,4] add(f32[1,4] prev_op3, f32[1,4] prev_op4) - op1 = f32[1,4] add(f32[1,4] prev_op4, f32[1,4] op0) - op2 = f32[1,4] add(f32[1,4] op0, f32[1,4] op1) - op3 = f32[1,4] add(f32[1,4] op1, f32[1,4] op2) - pinned_slice2 = f32[1,4] dynamic-slice(pinned_prev_param0, zero, zero), dynamic_slice_sizes={1,4} - op4 = f32[1,4] multiply(f32[1,4] pinned_slice2, f32[1,4] op3) - next_op0 = f32[1,4] add(f32[1,4] op3, f32[1,4] op4) - next_op1 = f32[1,4] add(f32[1,4] op4, f32[1,4] next_op0) - next_op2 = f32[1,4] add(f32[1,4] next_op0, f32[1,4] next_op1) - next_op3 = f32[1,4] add(f32[1,4] next_op1, f32[1,4] next_op2) - pinned_slice3 = f32[1,4] dynamic-slice(pinned_prev_param0, zero, zero), dynamic_slice_sizes={1,4} - next_op4 = f32[1,4] multiply(f32[1,4] pinned_slice3, f32[1,4] next_op3) - p = pred[] get-tuple-element(while_body_param), index=5 - ROOT root = tuple(pinned_prev_param0, next_param0, prev_prev_op3, prev_prev_op4, next_op4, p) - } - - ENTRY entry { - p0 = f32[28,4] parameter(0) - p1 = f32[1,4] parameter(1) - p2 = f32[1,4] parameter(2) - p3 = f32[1,4] parameter(3) - p4 = pred[] parameter(4) - copy = f32[1,4] copy(p3) - tuple = (f32[28,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) tuple(p0, p1, p2, p3, copy, p4) - while = (f32[28,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) while(tuple), condition=while_cond, body=while_body - ROOT root = f32[1,4] get-tuple-element(while), index=4 - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str)); - int64_t alternate_memory_size = 72; - ASSERT_OK_AND_ASSIGN(auto optimizer, CreateOptimizer(21, 27, module.get(), - alternate_memory_size)); - optimizer->Optimize(); - std::vector remaining_memory = optimizer->RemainingMemory(); - // We expect that pinned_prev_param0 would not get pinned due to negative - // savings: 32(uses) - 28 * 16(size) = -416 Time 0: 3 temporaries (16 B) + 1 - // pinned (4 B) - EXPECT_EQ(remaining_memory.at(0), alternate_memory_size - (3 * 16 + 4)); - EXPECT_EQ(optimizer->MaxAlternateMemoryUsed(), alternate_memory_size); -} - -TEST_F(MemoryBoundLoopOptimizerTest, OptimizerEndToEndWhileLoop) { - absl::string_view hlo_str = R"( -HloModule module, is_scheduled=true - -while_cond { - while_cond_param = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(while_cond_param), index=6 -} - -while_body { - while_body_param = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - prev_param0 = f32[1,4] get-tuple-element(while_body_param), index=0 - param0 = f32[1,4] get-tuple-element(while_body_param), index=1 - next_param0 = f32[1,4] get-tuple-element(while_body_param), index=2 - prev_prev_op3 = f32[1,4] get-tuple-element(while_body_param), index=3 - prev_prev_op4 = f32[1,4] get-tuple-element(while_body_param), index=4 - prev_op0 = f32[1,4] add(f32[1,4] prev_prev_op3, f32[1,4] prev_prev_op4) - prev_op1 = f32[1,4] add(f32[1,4] prev_prev_op4, f32[1,4] prev_op0) - prev_op2 = f32[1,4] add(f32[1,4] prev_op0, f32[1,4] prev_op1) - prev_op3 = f32[1,4] add(f32[1,4] prev_op1, f32[1,4] prev_op2) - prev_op4 = f32[1,4] multiply(f32[1,4] prev_param0, f32[1,4] prev_op3) - op0 = f32[1,4] add(f32[1,4] prev_op3, f32[1,4] prev_op4) - op1 = f32[1,4] add(f32[1,4] prev_op4, f32[1,4] op0) - op2 = f32[1,4] add(f32[1,4] op0, f32[1,4] op1) - op3 = f32[1,4] add(f32[1,4] op1, f32[1,4] op2) - op4 = f32[1,4] multiply(f32[1,4] param0, f32[1,4] op3) - next_op0 = f32[1,4] add(f32[1,4] op3, f32[1,4] op4) - next_op1 = f32[1,4] add(f32[1,4] op4, f32[1,4] next_op0) - next_op2 = f32[1,4] add(f32[1,4] next_op0, f32[1,4] next_op1) - next_op3 = f32[1,4] add(f32[1,4] next_op1, f32[1,4] next_op2) - next_op4 = f32[1,4] multiply(f32[1,4] next_param0, f32[1,4] next_op3) - p = pred[] get-tuple-element(while_body_param), index=6 - ROOT root = tuple(prev_param0, param0, next_param0, prev_prev_op3, prev_prev_op4, next_op4, p) -} - -ENTRY entry { - p0 = f32[1,4] parameter(0) - p1 = f32[1,4] parameter(1) - p2 = f32[1,4] parameter(2) - p3 = f32[1,4] parameter(3) - p4 = f32[1,4] parameter(4) - p5 = pred[] parameter(5) - copy = f32[1,4] copy(p4) - tuple = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) tuple(p0, p1, p2, p3, p4, copy, p5) - while = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) while(tuple), condition=while_cond, body=while_body - ROOT root = f32[1,4] get-tuple-element(while), index=5 -} - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str)); - - ASSERT_OK_AND_ASSIGN(auto preset_assignments, - RunMsa(module.get(), /*alternate_memory_size=*/512)); - - // We expect operand 0 of prev_op4, op4, and next_op4 to all be prefetches of - // same distance from the user. - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(auto hlo_live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - const HloInstruction* prev_copy_done = - FindInstruction(module.get(), "prev_op4")->operand(0); - const HloInstruction* copy_done = - FindInstruction(module.get(), "op4")->operand(0); - const HloInstruction* next_copy_done = - FindInstruction(module.get(), "next_op4")->operand(0); - ASSERT_EQ(prev_copy_done->opcode(), HloOpcode::kCopyDone); - ASSERT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); - ASSERT_EQ(next_copy_done->opcode(), HloOpcode::kCopyDone); - EXPECT_EQ(prev_copy_done->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(copy_done->shape().layout().memory_space(), kAlternateMemorySpace); - EXPECT_EQ(next_copy_done->shape().layout().memory_space(), - kAlternateMemorySpace); - auto prefetch_distance = [&](const HloInstruction* copy_done) { - return hlo_live_range->instruction_schedule().at(copy_done) - - hlo_live_range->instruction_schedule().at(copy_done->operand(0)); - }; - EXPECT_EQ(prefetch_distance(prev_copy_done), prefetch_distance(copy_done)); - EXPECT_EQ(prefetch_distance(next_copy_done), prefetch_distance(copy_done)); -} - -TEST_F(MemoryBoundLoopOptimizerTest, OptimizerEndToEndNestedWhileLoopBug) { - absl::string_view hlo_str = R"( -HloModule module, is_scheduled=true - -prev_while_cond { - prev_while_cond_param = (f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(prev_while_cond_param), index=1 -} - -prev_while_body { - prev_while_body_param = (f32[1,4], pred[]) parameter(0) - prev_while_body_gte = f32[1,4] get-tuple-element(prev_while_body_param), index=0 - prev_while_body_pred = pred[] get-tuple-element(prev_while_body_param), index=1 - prev_while_body_op = f32[1,4] negate(prev_while_body_gte) - ROOT prev_while_body_root = (f32[1,4], pred[]) tuple(prev_while_body_op, prev_while_body_pred) -} - -current_while_cond { - current_while_cond_param = (f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(current_while_cond_param), index=1 -} - -current_while_body { - current_while_body_param = (f32[1,4], pred[]) parameter(0) - current_while_body_gte = f32[1,4] get-tuple-element(current_while_body_param), index=0 - current_while_body_pred = pred[] get-tuple-element(current_while_body_param), index=1 - current_while_body_op = f32[1,4] negate(current_while_body_gte) - ROOT current_while_body_root = (f32[1,4], pred[]) tuple(current_while_body_op, current_while_body_pred) -} - -next_while_cond { - next_while_cond_param = (f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(next_while_cond_param), index=1 -} - -next_while_body { - next_while_body_param = (f32[1,4], pred[]) parameter(0) - next_while_body_gte = f32[1,4] get-tuple-element(next_while_body_param), index=0 - next_while_body_pred = pred[] get-tuple-element(next_while_body_param), index=1 - next_while_body_op = f32[1,4] negate(next_while_body_gte) - ROOT next_while_body_root = (f32[1,4], pred[]) tuple(next_while_body_op, next_while_body_pred) -} - -while_cond { - while_cond_param = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - ROOT p = pred[] get-tuple-element(while_cond_param), index=6 -} - -while_body { - while_body_param = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) parameter(0) - prev_param0 = f32[1,4] get-tuple-element(while_body_param), index=0 - param0 = f32[1,4] get-tuple-element(while_body_param), index=1 - next_param0 = f32[1,4] get-tuple-element(while_body_param), index=2 - prev_prev_op3 = f32[1,4] get-tuple-element(while_body_param), index=3 - prev_prev_op4 = f32[1,4] get-tuple-element(while_body_param), index=4 - while_pred = pred[] get-tuple-element(while_body_param), index=6 - prev_op0 = f32[1,4] add(f32[1,4] prev_prev_op3, f32[1,4] prev_prev_op4) - prev_op1 = f32[1,4] add(f32[1,4] prev_prev_op4, f32[1,4] prev_op0) - prev_op2 = f32[1,4] add(f32[1,4] prev_op0, f32[1,4] prev_op1) - prev_op3 = f32[1,4] add(f32[1,4] prev_op1, f32[1,4] prev_op2) - prev_tuple = (f32[1,4], pred[]) tuple(prev_op3, while_pred) - prev_while = (f32[1,4], pred[]) while(prev_tuple), condition=prev_while_cond, body=prev_while_body - prev_gte = f32[1,4] get-tuple-element(prev_while), index=0 - prev_op4 = f32[1,4] multiply(f32[1,4] prev_param0, f32[1,4] prev_gte) - op0 = f32[1,4] add(f32[1,4] prev_op3, f32[1,4] prev_op4) - op1 = f32[1,4] add(f32[1,4] prev_op4, f32[1,4] op0) - op2 = f32[1,4] add(f32[1,4] op0, f32[1,4] op1) - op3 = f32[1,4] add(f32[1,4] op1, f32[1,4] op2) - current_tuple = (f32[1,4], pred[]) tuple(op3, while_pred) - current_while = (f32[1,4], pred[]) while(current_tuple), condition=current_while_cond, body=current_while_body - current_gte = f32[1,4] get-tuple-element(current_while), index=0 - op4 = f32[1,4] multiply(f32[1,4] param0, f32[1,4] current_gte) - next_op0 = f32[1,4] add(f32[1,4] op3, f32[1,4] op4) - next_op1 = f32[1,4] add(f32[1,4] op4, f32[1,4] next_op0) - next_op2 = f32[1,4] add(f32[1,4] next_op0, f32[1,4] next_op1) - next_op3 = f32[1,4] add(f32[1,4] next_op1, f32[1,4] next_op2) - next_tuple = (f32[1,4], pred[]) tuple(next_op3, while_pred) - next_while = (f32[1,4], pred[]) while(next_tuple), condition=next_while_cond, body=next_while_body - next_gte = f32[1,4] get-tuple-element(next_while), index=0 - next_op4 = f32[1,4] multiply(f32[1,4] next_param0, f32[1,4] next_gte) - ROOT root = tuple(prev_param0, param0, next_param0, prev_prev_op3, prev_prev_op4, next_op4, while_pred) -} - -ENTRY entry { - p0 = f32[1,4] parameter(0) - p1 = f32[1,4] parameter(1) - p2 = f32[1,4] parameter(2) - p3 = f32[1,4] parameter(3) - p4 = f32[1,4] parameter(4) - p5 = pred[] parameter(5) - copy = f32[1,4] copy(p4) - tuple = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) tuple(p0, p1, p2, p3, p4, copy, p5) - while = (f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], f32[1,4], pred[]) while(tuple), condition=while_cond, body=while_body - ROOT root = f32[1,4] get-tuple-element(while), index=5 -} - )"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str)); - - ASSERT_OK_AND_ASSIGN(auto preset_assignments, - RunMsa(module.get(), /*alternate_memory_size=*/512)); -} - -TEST_F(MemoryBoundLoopOptimizerTest, TuplePositionAfterUseDoesNotCrash) { - absl::string_view hlo_string = R"( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[1,4] parameter(0) - p1 = f32[1,4] parameter(1) - p2 = f32[1,4] parameter(2) - p3 = f32[1,4] parameter(3) - - // Iteration 0 - c0 = f32[1,4] tanh(p0) - t0 = (f32[1,4], f32[1,4]) tuple(c0, p0) - op1_0 = f32[1,4] tanh(c0) - op2_0 = f32[1,4] add(op1_0, op1_0) - - // Iteration 1 - c1 = f32[1,4] tanh(p1) - t1 = (f32[1,4], f32[1,4]) tuple(c1, p1) - op1_1 = f32[1,4] tanh(c1) - op2_1 = f32[1,4] add(op1_1, op1_1) - - // Iteration 2 - c2 = f32[1,4] tanh(p2) - t2 = (f32[1,4], f32[1,4]) tuple(c2, p2) - op1_2 = f32[1,4] tanh(c2) - op2_2 = f32[1,4] add(op1_2, op1_2) - - ROOT root = tuple(t0, t1, t2, op2_0, op2_1, op2_2) - })"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(auto preset_assignments, RunMsa(module.get())); -} - -} // namespace -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc deleted file mode 100644 index a4bd5aa6a63452..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc +++ /dev/null @@ -1,1713 +0,0 @@ -/* Copyright 2019 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/memory_space_assignment.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/base/casts.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/log/vlog_is_on.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_format.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.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/heap_simulator/heap_simulator.h" -#include "xla/service/hlo.pb.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/algorithm.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/service/memory_space_assignment/simulator.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -#include "xla/util.h" -#include "xla/xla_data.pb.h" -#include "tsl/platform/casts.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -// Insert an instruction to the schedule, and make sure its dependencies -// (operands) are already in the schedule. If not, insert these operands -// before the instruction. -void InsertInstructionAndEnsureOperandsInserted( - HloInstruction* new_instruction, HloInstructionSequence* new_sequence, - absl::flat_hash_set* inserted_instructions) { - if (!inserted_instructions->insert(new_instruction).second) { - VLOG(1) << "Already inserted: " << new_instruction->ToString(); - } else { - for (HloInstruction* operand : new_instruction->operands()) { - InsertInstructionAndEnsureOperandsInserted(operand, new_sequence, - inserted_instructions); - } - VLOG(1) << "Inserting: " << new_instruction->ToShortString(); - new_sequence->push_back(new_instruction); - } -} - -std::string InstructionScheduleToString(const HloLiveRange& hlo_live_range) { - const absl::flat_hash_map& - instruction_schedule = hlo_live_range.instruction_schedule(); - std::vector> instructions; - instructions.reserve(instruction_schedule.size()); - for (const auto& instruction : instruction_schedule) { - instructions.push_back({instruction.second, instruction.first}); - } - std::string instruction_schedule_str = "\n"; - absl::c_sort(instructions); - for (auto& instruction : instructions) { - absl::StrAppend(&instruction_schedule_str, - "LogicalTime: ", instruction.first, " ", - instruction.second->ToString(), "\n"); - } - return instruction_schedule_str; -} - -void EnsureParentAllocationIsAvailableForCopy(CopyAllocation* copy_allocation) { - Allocation& parent_allocation = copy_allocation->mutable_prev_allocation(); - parent_allocation.Extend(copy_allocation->copy_done_schedule_before()); - if (parent_allocation.is_copy_allocation()) { - auto parent_copy_allocation = - absl::down_cast(&parent_allocation); - parent_copy_allocation->set_copy_done_schedule_before( - std::min(parent_copy_allocation->copy_done_schedule_before(), - copy_allocation->start_time())); - parent_copy_allocation->set_copy_start_schedule_after( - std::min(parent_copy_allocation->copy_start_schedule_after(), - parent_copy_allocation->copy_done_schedule_before() - 1)); - } -} - -void MakeCopyAllocationJitForSingleUse(CopyAllocation* copy_allocation, - int64_t use_time) { - copy_allocation->set_start_time(use_time - 1); - copy_allocation->set_copy_start_schedule_after(use_time - 1); - copy_allocation->set_end_time(use_time); - copy_allocation->set_copy_done_schedule_before(use_time); - EnsureParentAllocationIsAvailableForCopy(copy_allocation); -} - -int64_t GetUseTime(const HloUse& use, const HloLiveRange& hlo_live_range) { - return hlo_live_range.instruction_schedule().at(use.instruction); -} - -void ProcessPrefetchesToAlternateMemory(AllocationSequence& allocations, - const HloLiveRange& hlo_live_range) { - std::vector allocations_in_raw_pointers = - GetAllocationSequenceInRawPointers(allocations); - for (auto allocation : allocations_in_raw_pointers) { - if (allocation->is_copy_allocation() && allocation->is_in_alternate_mem() && - !allocation->uses().empty()) { - CopyAllocation* prefetch = absl::down_cast(allocation); - std::vector uses = prefetch->uses(); // Create a copy of uses. - prefetch->clear_uses(); // Clear old uses. - // For every prefetch, update prefetch to serve earliest use just in time. - prefetch->AddUse(uses[0]); - MakeCopyAllocationJitForSingleUse(prefetch, - GetUseTime(uses[0], hlo_live_range)); - // For every use after the first use, create a new prefetch from the same - // parent allocation. - for (size_t use_index = 1; use_index < uses.size(); ++use_index) { - const HloUse& use = uses[use_index]; - int64_t use_time = GetUseTime(use, hlo_live_range); - auto jit_single_use_prefetch = std::make_unique( - prefetch->mutable_prev_allocation(), MemorySpace::kAlternate, - prefetch->chunk(), use_time - 1, use_time, use_time); - jit_single_use_prefetch->set_copy_start_schedule_after(use_time - 1); - jit_single_use_prefetch->AddUse(use); - EnsureParentAllocationIsAvailableForCopy(jit_single_use_prefetch.get()); - allocations.push_back(std::move(jit_single_use_prefetch)); - } - } - } -} - -void MakeEvictionImmediate(CopyAllocation* eviction) { - const Allocation& parent_allocation = eviction->prev_allocation(); - eviction->set_start_time(parent_allocation.start_time()); - eviction->set_copy_start_schedule_after(parent_allocation.start_time()); - eviction->set_copy_done_schedule_before(parent_allocation.start_time() + 1); - eviction->Extend(parent_allocation.start_time() + 1); -} - -absl::flat_hash_map GetEvictionsMap( - std::vector& allocations) { - absl::flat_hash_map evictions_map; - for (auto& allocation : allocations) { - if (allocation->is_copy_allocation() && allocation->is_in_default_mem()) { - auto eviction = absl::down_cast(allocation); - Allocation& parent_allocation = eviction->mutable_prev_allocation(); - if (!parent_allocation.is_copy_allocation()) { - evictions_map[&parent_allocation] = eviction; - } - } - } - return evictions_map; -} - -void ProcessBuffersProducedInAlternateMemory( - AllocationSequence& allocations, const HloLiveRange& hlo_live_range) { - std::vector allocations_in_raw_pointers = - GetAllocationSequenceInRawPointers(allocations); - // For all parent allocations produced in alternate memory, create a map from - // parent allocation -> eviction. - absl::flat_hash_map evictions_map = - GetEvictionsMap(allocations_in_raw_pointers); - // Make all such evictions immediate. - for (auto& [_, eviction] : evictions_map) { - MakeEvictionImmediate(eviction); - } - if (VLOG_IS_ON(2)) { - LOG(INFO) << "AllocationSequence after making spills immediate spills\n"; - XLA_LOG_LINES(INFO, AllocationSequenceToString(allocations, true)); - } - // Process all buffers produced in the alternate memory: - // 1. Make the buffer short lived. - // 2. Service immediate use if any. - // 3. If buffer is also used later get or create an immediate eviction. - // 4. For every later use prefetch just in time from the eviction. - for (auto allocation : allocations_in_raw_pointers) { - if (!allocation->is_copy_allocation() && - allocation->is_in_alternate_mem()) { - std::vector uses = allocation->uses(); // Create a copy of uses. - allocation->clear_uses(); // Clear old uses. - // Make buffer short lived. - allocation->set_end_time(allocation->start_time() + 1); - for (const HloUse& use : uses) { - int64_t use_time = GetUseTime(use, hlo_live_range); - if (allocation->start_time() + 1 == use_time) { - allocation->AddUse(use); - continue; - } - if (!evictions_map.contains(allocation)) { - auto eviction_unique_ptr = std::make_unique( - *allocation, MemorySpace::kDefault, std::nullopt, - allocation->start_time(), allocation->start_time() + 1, - allocation->start_time() + 1); - eviction_unique_ptr->set_copy_start_schedule_after( - allocation->start_time()); - evictions_map[allocation] = eviction_unique_ptr.get(); - allocations.push_back(std::move(eviction_unique_ptr)); - } - CopyAllocation* eviction = evictions_map[allocation]; - auto jit_single_use_prefetch = std::make_unique( - *eviction, MemorySpace::kAlternate, allocation->chunk(), - use_time - 1, use_time, use_time); - jit_single_use_prefetch->set_copy_start_schedule_after(use_time - 1); - jit_single_use_prefetch->AddUse(use); - EnsureParentAllocationIsAvailableForCopy(jit_single_use_prefetch.get()); - allocations.push_back(std::move(jit_single_use_prefetch)); - } - } - } -} - -void TransformAllocationSequenceToSpill(AllocationSequence& allocations, - const HloLiveRange& hlo_live_range) { - if (VLOG_IS_ON(2)) { - LOG(INFO) << "InstructionSchedule before transform\n"; - XLA_LOG_LINES(INFO, InstructionScheduleToString(hlo_live_range)); - LOG(INFO) << "AllocationSequence before transform\n"; - XLA_LOG_LINES(INFO, AllocationSequenceToString(allocations, true)); - } - ProcessPrefetchesToAlternateMemory(allocations, hlo_live_range); - if (VLOG_IS_ON(2)) { - LOG(INFO) << "AllocationSequence after processing prefetches\n"; - XLA_LOG_LINES(INFO, AllocationSequenceToString(allocations, true)); - } - ProcessBuffersProducedInAlternateMemory(allocations, hlo_live_range); - if (VLOG_IS_ON(2)) { - VLOG(2) << "AllocationSequence after processing buffers produced in kAlt\n"; - XLA_LOG_LINES(INFO, AllocationSequenceToString(allocations, true)); - } - SortAllocationSequence(allocations); -} - -} // namespace - -absl::StatusOr -MemorySpaceAssignment::CalculateAsyncCopyStats( - const HloDataflowAnalysis& dataflow_analysis) const { - AsyncCopyStats stats; - int64_t current_copies = 0; - for (const HloComputation* computation : - module_->MakeNonfusionComputations()) { - for (HloInstruction* instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::kCopyStart || - (instruction->opcode() == HloOpcode::kAsyncStart && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice)) { - current_copies++; - } else if (instruction->opcode() == HloOpcode::kCopyDone || - (instruction->opcode() == HloOpcode::kAsyncDone && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice)) { - current_copies--; - int64_t size = - options_.size_fn(dataflow_analysis.GetUniqueValueAt(instruction)); - if (instruction->shape().layout().memory_space() == - options_.alternate_memory_space) { - ++stats.num_prefetches; - stats.prefetch_bytes += size; - if (instruction->opcode() == HloOpcode::kAsyncDone && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice) { - ++stats.num_sliced_prefetch_slices; - } - } else { - ++stats.num_evictions; - stats.eviction_bytes += size; - } - } else if (instruction->IsCustomCall(kConcatBitcastCustomCall)) { - ++stats.num_sliced_prefetches; - } - stats.max_outstanding_async_copies = - std::max(stats.max_outstanding_async_copies, current_copies); - } - } - return stats; -} - -/*static*/ absl::StatusOr> -MemorySpaceAssignment::Run(HloModule* module, - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis, - const AliasInfo* alias_info, - const Options& options) { - CHECK(module->has_schedule()); - if (VLOG_IS_ON(3)) { - LOG(INFO) << "memory_space_assignment_options::Options:\n"; - XLA_LOG_LINES(INFO, options.ToString()); - } - if (VLOG_IS_ON(10)) { - LOG(INFO) << "Module before memory space assignment: "; - XLA_LOG_LINES(INFO, module->ToString()); - LOG(INFO) << "Schedule: " << module->schedule().ToString(); - } - MemorySpaceAssignment memory_space_assignment(module, alias_info, options, - hlo_live_range); - - return memory_space_assignment.RunMemorySpaceAssignment(hlo_live_range, - alias_analysis); -} - -absl::Status MemorySpaceAssignment::VerifyAllocations() const { - BufferIntervalTree interval_tree; - // Checks the chunks that overlap with a given allocation in time do not - // overlap with the allocation's chunk in the memory range. If they do, we - // throw an error, otherwise we add the allocation's chunk to the interval - // tree and return an OK status. - auto add_allocation_and_verify = - [&](const Allocation* allocation) -> absl::Status { - for (const HeapSimulator::Chunk& overlapping_chunk : - interval_tree.ChunksOverlappingInTime(allocation->start_time(), - allocation->end_time() - 1)) { - CHECK(!allocation->chunk().OverlapsWith(overlapping_chunk)) - << "Chunks are overlapping at Allocation level (before fixing the " - "schedule): " - << allocation->ToString() - << " overlaps with allocated chunk: " << overlapping_chunk.ToString(); - } - interval_tree.Add(allocation->start_time(), allocation->end_time() - 1, - allocation->chunk()); - return absl::OkStatus(); - }; - // Verify that all alternate memory allocations are free of overlapping - // Allocations in time and space, and add them to interval_tree one by one. - for (const auto& allocation : allocations_) { - if (allocation->memory_space() == MemorySpace::kAlternate && - !allocation->is_mirrored_allocation()) { - ABSL_RETURN_IF_ERROR(add_allocation_and_verify(allocation.get())); - } - } - return absl::OkStatus(); -} - -absl::StatusOr> -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_.IsOpSpanExposureEnabled()) - << "TODO(b/167392593): Support split shapes for window " - "prefetches."; - } - ABSL_RETURN_IF_ERROR(FindAllocationSequence(hlo_live_range, alias_analysis)); - - std::optional runtime_simulator = std::nullopt; - if (options_.cost_analysis) { - runtime_simulator.emplace(options_.cost_analysis, - options_.alternate_memory_space); - if (VLOG_IS_ON(1)) { - float estimated_time = - runtime_simulator->SimulateElapsedTimeWithoutAsyncCopyLikes( - hlo_live_range, allocations_); - LOG(INFO) << "Estimated elapsed time without async copies (sec): " - << estimated_time; - } - } - - ABSL_RETURN_IF_ERROR(Process(hlo_live_range, alias_analysis)); - if (options_.verify) { - ABSL_RETURN_IF_ERROR(VerifyAllocations()); - } - - // DEBUG_LOG_ALLOCATIONS_AT - // - // Uncomment the following to log the alternate memory allocations that MSA - // made at a given schedule time. - // - // AllocationSequenceDebugging::LogAltMemAllocationsAt( - // allocations_, /*time*/1); - ScheduleAsynchronousCopies(); - ABSL_RETURN_IF_ERROR(SimplifyGraph()); - ABSL_RETURN_IF_ERROR(SetSchedule()); - ABSL_ASSIGN_OR_RETURN(auto alias, HloAliasAnalysis::Run(module_, alias_info_)); - ABSL_RETURN_IF_ERROR(ExportAndColorBuffers(*alias)); - std::vector alt_mem_bytes_occupied; - // alt_mem_bytes_occupied is used for logging in the RuntimeSimulator below. - // We only populate it in VerifyAndExportHeapSimulatorTrace if the - // RuntimeSimulator is present. - ABSL_RETURN_IF_ERROR(VerifyAndExportHeapSimulatorTrace( - *alias, - runtime_simulator.has_value() ? &alt_mem_bytes_occupied : nullptr)); - if (VLOG_IS_ON(2) && runtime_simulator.has_value()) { - float estimated_time = runtime_simulator->SimulateElapsedTime( - module_, *alias, allocations_, &alt_mem_bytes_occupied); - LOG(INFO) << "Estimated elapsed time with async copies (sec): " - << estimated_time; - } - if (VLOG_IS_ON(3)) { - LOG(INFO) << "Module after memory space assignment: "; - XLA_LOG_LINES(INFO, module_->ToString()); - } - CHECK_OK(module_->schedule().Verify()); - if (VLOG_IS_ON(1)) { - ABSL_ASSIGN_OR_RETURN(AsyncCopyStats stats, - CalculateAsyncCopyStats(alias->dataflow_analysis())); - LOG(INFO) << "Maximum number of outstanding async copies/slices: " - << stats.max_outstanding_async_copies; - LOG(INFO) << "Number of prefetches: " << stats.num_prefetches - << ", in bytes: " << stats.prefetch_bytes; - LOG(INFO) << "Number of sliced prefetches: " << stats.num_sliced_prefetches - << ", consuming number of slices: " - << stats.num_sliced_prefetch_slices; - LOG(INFO) << "Number of evictions: " << stats.num_evictions - << ", in bytes: " << stats.eviction_bytes; - } - - return std::move(preset_assignments_); -} - -absl::Status MemorySpaceAssignment::FindAllocationSequence( - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - auto algorithm = std::make_unique(module_, &allocations_, - options_, alias_analysis, - alias_info_, hlo_live_range); - - HeapSimulator::Options heap_simulator_options; - heap_simulator_options.may_reuse_operand_buffers = false; - heap_simulator_options.alloc_constants = true; - return HeapSimulator::Run(std::move(algorithm), *module_, module_->schedule(), - alias_analysis, alias_info_, &options_.size_fn, - heap_simulator_options) - .status(); -} - -MemorySpaceAssignment::ScopedMemorySource -MemorySpaceAssignment::ScopedMemorySource::ForInstruction( - HloInstruction* instruction) { - return ScopedMemorySource{false, instruction}; -} - -MemorySpaceAssignment::ScopedMemorySource -MemorySpaceAssignment::ScopedMemorySource::ForPostModule() { - return ScopedMemorySource{true, nullptr}; -} - -std::string MemorySpaceAssignment::ScopedMemorySource::ToString() const { - if (is_post_module) { - return ""; - } - return std::string(instruction->name()); -} - -absl::Status MemorySpaceAssignment::Process( - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis) { - VLOG(1) << "Processing assigned buffers..."; - // Since some parent allocations may not be needed (e.g. when they don't have - // any uses and if there is no other (non-parent) allocation that depends on - // it, before we process the allocations, mark all allocations that are - // needed. - absl::flat_hash_set needed_allocations; - if (options_.always_spill_to_default_memory) { - TransformAllocationSequenceToSpill(allocations_, hlo_live_range); - } - for (auto& allocation : allocations_) { - allocation->MarkIfNeeded(needed_allocations); - } - // Insert CopyStart/CopyDone and SliceStart/SliceDone pairs. - for (auto& allocation : allocations_) { - VLOG(3) << "Processing: " << allocation->ToString(); - if (!needed_allocations.contains(allocation.get())) { - VLOG(3) << "Allocation not needed."; - continue; - } - ABSL_RETURN_IF_ERROR(allocation->Process(options_.bitcast_split_fn, - hlo_live_range, alias_analysis)); - // Add the offset and size of the allocation in the alternate memory to - // the output map. - if (allocation->is_scoped_allocation()) { - CHECK(allocation->memory_space() == MemorySpace::kAlternate); - ScopedAllocation* scoped_allocation = - static_cast(allocation.get()); - if (scoped_allocation->is_post_module()) { - scoped_memory_assignments_.emplace_back( - ScopedMemorySource::ForPostModule(), allocation->chunk()); - } else { - scoped_memory_assignments_.emplace_back( - ScopedMemorySource::ForInstruction( - scoped_allocation->defining_position().instruction), - allocation->chunk()); - } - alternate_memory_size_ = - std::max(alternate_memory_size_, allocation->chunk().chunk_end()); - } else if (allocation->memory_space() == MemorySpace::kAlternate) { - if (allocation->is_sliced_copy_allocation()) { - // Add slices - const SlicedCopyAllocation& sliced_copy_allocation = - *static_cast(allocation.get()); - for (const SlicedCopyAllocation::SliceDetail& details : - sliced_copy_allocation.slice_details_sorted_by_start_time()) { - alternate_memory_assignments_.push_back( - {{details.copy_done, {}}, details.slice_decision.chunk}); - alternate_memory_size_ = std::max( - alternate_memory_size_, details.slice_decision.chunk.chunk_end()); - } - CHECK( - !sliced_copy_allocation.cross_program_prefetch_index().has_value()); - } - alternate_memory_assignments_.emplace_back( - allocation->defining_position(), allocation->chunk()); - alternate_memory_size_ = - std::max(alternate_memory_size_, allocation->chunk().chunk_end()); - if (allocation->split_shape().has_value()) { - auto result = split_map_.insert({allocation->defining_position(), - &allocation->split_shape()->layout()}); - if (!result.second) { - CHECK_EQ(*result.first->second, - allocation->split_shape().value().layout()); - } - } - if (allocation->cross_program_prefetch_index().has_value()) { - ABSL_RETURN_IF_ERROR(module_->SetCrossProgramPrefetchOffset( - *allocation->cross_program_prefetch_index(), - allocation->chunk().offset)); - } - } - } - - // Post-process allocations. This is only used for parent allocations where we - // update the body root with a reference to the buffer in default memory - // space. - absl::flat_hash_set seen_pinned_positions; - for (auto& allocation : allocations_) { - if (needed_allocations.contains(allocation.get())) { - VLOG(3) << "Post-Processing: " << allocation->ToString(); - ABSL_RETURN_IF_ERROR(allocation->PostProcess()); - if (allocation->is_pinned_allocation() && - !allocation->is_scoped_allocation()) { - auto [it, inserted] = - seen_pinned_positions.insert(allocation->defining_position()); - TF_RET_CHECK(inserted) - << "Multiple pinned allocations defined for position " - << allocation->defining_position().ToString(); - } - } - } - return absl::OkStatus(); -} - -absl::Status MemorySpaceAssignment::ExportAndColorBuffers( - const HloAliasAnalysis& alias_analysis) { - VLOG(1) << "Exporting buffers..."; - absl::flat_hash_map seen_buffer_offsets; - VLOG(3) << "Exported alternate memory allocations:"; - for (const auto& position_and_chunk : alternate_memory_assignments_) { - const HloPosition& defining_position = position_and_chunk.first; - const HeapSimulator::Chunk& chunk = position_and_chunk.second; - const HloBuffer& buffer = alias_analysis.GetUniqueBufferAt( - defining_position.instruction, defining_position.index); - auto seen_buffer_offset_it = seen_buffer_offsets.find(buffer.id()); - if (seen_buffer_offset_it != seen_buffer_offsets.end()) { - CHECK_EQ(chunk.offset, seen_buffer_offset_it->second) - << "Mismatch in offset for positions that map to the same value: " - << buffer.ToString() << ", pos: " << defining_position.ToString(); - } else { - VLOG(3) << " [" << chunk.offset << ", " << chunk.size - << "] : " << defining_position.ToString() << " (" - << buffer.ToString() << ")"; - preset_assignments_->add_chunk(defining_position, chunk); - seen_buffer_offsets[buffer.id()] = chunk.offset; - } - } - - VLOG(3) << "Exported scoped allocations in alternate memory:"; - for (const auto& instruction_and_chunk : scoped_memory_assignments_) { - ScopedMemorySource scoped_memory_source = instruction_and_chunk.first; - const HeapSimulator::Chunk& chunk = instruction_and_chunk.second; - if (scoped_memory_source.is_post_module) { - preset_assignments_->set_post_module_scoped_alternate_memory_chunk(chunk); - } else { - preset_assignments_->add_scoped_allocation_chunk( - scoped_memory_source.instruction, chunk); - } - VLOG(3) << " [" << chunk.offset << ", " << chunk.size - << "] : " << scoped_memory_source.ToString(); - } - - if (!preset_assignments_->chunks().empty() || - !preset_assignments_->scoped_allocation_chunks().empty()) { - preset_assignments_ - ->assignment_information_for_space(options_.alternate_memory_space) - ->size = alternate_memory_size_; - } - - VLOG(3) << "Exported alternate memory sizes:"; - for (auto& pair : preset_assignments_->assignment_informations()) { - VLOG(3) << " space: " << pair.first << ", size: " << pair.second.size; - } - - VLOG(1) << "Coloring buffers..."; - // Color the pending positions and all of their aliased buffers. - for (const auto& defining_position_and_chunk : - preset_assignments_->chunks()) { - const HloPosition& defining_position = defining_position_and_chunk.first; - auto split_result = split_map_.find(defining_position); - for (auto& buffer : alias_analysis.ComputeBuffersAt( - defining_position.instruction, defining_position.index)) { - for (auto& value : buffer->values()) { - for (auto& position : value->positions()) { - VLOG(4) << "Coloring " << position.ToString(); - Shape* shape = ShapeUtil::GetMutableSubshape( - position.instruction->mutable_shape(), position.index); - CHECK(shape->IsArray()) << "Coloring a shape that is not an array: " - << position.ToString(); - shape->mutable_layout()->set_memory_space( - options_.alternate_memory_space); - // For asynchronous dynamic slice and dynamic update-slice start - // instructions, propagate the alternate memory space coloring - // directly to the wrapped instruction's layout so downstream code - // generation sees the correct memory space. - if (position.instruction->opcode() == HloOpcode::kAsyncStart && - position.index == ShapeIndex({1}) && - position.instruction->async_wrapped_instruction() != nullptr) { - HloOpcode wrapped_op = position.instruction->async_wrapped_opcode(); - if (wrapped_op == HloOpcode::kDynamicUpdateSlice || - wrapped_op == HloOpcode::kDynamicSlice) { - Shape* wrapped_shape = - position.instruction->async_wrapped_instruction() - ->mutable_shape(); - if (wrapped_shape->has_layout()) { - wrapped_shape->mutable_layout()->set_memory_space( - options_.alternate_memory_space); - } - } - } - if (split_result != split_map_.end()) { - CHECK_EQ(shape->layout().split_configs().size(), 0); - shape->mutable_layout()->add_split_configs( - split_result->second->split_configs(0)); - } - } - } - } - } - - return absl::OkStatus(); -} - -void MemorySpaceAssignment::RemoveAlternateMemoryAssignments( - const absl::flat_hash_set& instructions) { - for (auto it = alternate_memory_assignments_.begin(); - it != alternate_memory_assignments_.end();) { - const HloInstruction* instruction = it->first.instruction; - if (!instructions.contains(instruction)) { - ++it; - continue; - } - - VLOG(3) << "Removing instruction from alternate memory assignments."; - if (std::next(it) == alternate_memory_assignments_.end()) { - alternate_memory_assignments_.pop_back(); - break; - } - - // Swap the removed position and chunk with the back and pop back. - *it = alternate_memory_assignments_.back(); - alternate_memory_assignments_.pop_back(); - } -} - -void MemorySpaceAssignment::RemoveScopedMemoryAssignments( - const absl::flat_hash_set& instructions) { - for (auto it = scoped_memory_assignments_.begin(); - it != scoped_memory_assignments_.end();) { - const HloInstruction* instruction = it->first.instruction; - if (!instructions.contains(instruction)) { - ++it; - continue; - } - - VLOG(3) << "Removing instruction from alternate memory assignments."; - if (std::next(it) == scoped_memory_assignments_.end()) { - scoped_memory_assignments_.pop_back(); - break; - } - - // Swap the removed position and chunk with the back and pop back. - *it = scoped_memory_assignments_.back(); - scoped_memory_assignments_.pop_back(); - } -} - -namespace { - -// Removes the given computation, and adds computations to the worklist that -// will be orphaned as a result. Removed instructions are added to -// removed_instructions. -absl::Status ProcessDeadComputation( - HloModule* module, HloComputation* computation, - absl::flat_hash_set& removed_instructions, - std::vector& worklist) { - for (HloInstruction* instruction : computation->instructions()) { - removed_instructions.insert(instruction); - for (HloComputation* called_computation : - instruction->called_computations()) { - // If computation is cloned, the called computation might still have other - // callers, so not necessarily dead. - if (called_computation->caller_computations().size() == 1) { - worklist.push_back(called_computation); - } - } - } - VLOG(2) << "Removing dead computation: " << computation->name(); - ABSL_RETURN_IF_ERROR(module->RemoveEmbeddedComputation(computation)); - return absl::OkStatus(); -} - -// Removes dead fusion computations from the module and populates -// removed_instructions with the instructions that were removed. -absl::Status CleanupDeadFusionComputations( - HloModule* module, - absl::flat_hash_set& removed_instructions) { - std::vector worklist; - for (HloComputation* computation : module->computations()) { - if (!computation->IsEntryComputation() && - computation->caller_computations().empty()) { - worklist.push_back(computation); - } - } - while (!worklist.empty()) { - HloComputation* computation = worklist.back(); - worklist.pop_back(); - ABSL_RETURN_IF_ERROR(ProcessDeadComputation(module, computation, - removed_instructions, worklist)); - } - return absl::OkStatus(); -} - -} // namespace - -absl::Status MemorySpaceAssignment::SimplifyGraph() { - VLOG(1) << "Simplifying graph..."; - - // Preprocess flattened_instructions_ for quick index lookup. - absl::flat_hash_map - instruction_to_flattened_instructions_idx; - for (int64_t i = 0; i < flattened_instructions_.size(); ++i) { - if (flattened_instructions_[i] == nullptr) { - continue; - } - instruction_to_flattened_instructions_idx[flattened_instructions_[i]] = i; - } - - absl::flat_hash_set removed_instructions; - - for (HloComputation* computation : module_->MakeNonfusionComputations()) { - // Parallel computations aren't in the schedule and don't need to be - // modified. - if (!computations_in_schedule_.contains(computation)) { - VLOG(4) << "Not simplifying " << computation->name() - << " because it's not in the schedule."; - continue; - } - // Drop control dependencies. Since the computation is already scheduled, we - // don't need control dependencies anymore, and having control - // predecessors/successors prevents us from removing instructions without - // users (HloComputation::IsSafelyRemovable returns false if there are - // control dependencies). - for (HloInstruction* instruction : - computation->MakeInstructionPostOrder()) { - ABSL_RETURN_IF_ERROR(instruction->DropAllControlDeps()); - } - // We perform limited DCE and forward the tuple operand in patterns like - // GetTupleElement(Tuple(a, b), 0). This is mostly because memory space - // assignment is ran late in compilation (after DCE and arithmetic - // simplification passes) and we don't want to generate redundant code. Run - // to fixed point. - bool computation_modified = true; - while (computation_modified) { - computation_modified = false; - VLOG(4) << "Running simplify graph loop over " << computation->name(); - for (HloInstruction* instruction : - computation->MakeInstructionPostOrder()) { - if (computation->IsSafelyRemovable(instruction) && - instruction->IsDead() && !instruction->HasSideEffect() && - instruction->opcode() != HloOpcode::kCopyStart && - instruction->opcode() != HloOpcode::kCopyDone) { - VLOG(4) << "Instruction removed: " << instruction->ToString(); - removed_instructions.insert(instruction); - // Instead of deleting the instruction from the schedule, replace it - // with a nullptr. This is needed because SetSchedule relies on the - // logical time that is the index into flattened_instructions_ for - // scheduling asynchronous copies. - if (instruction_to_flattened_instructions_idx.contains(instruction)) { - flattened_instructions_ - [instruction_to_flattened_instructions_idx[instruction]] = - nullptr; - } - ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(instruction)); - computation_modified = true; - } else if (instruction->opcode() == HloOpcode::kGetTupleElement) { - HloInstruction* operand = instruction->mutable_operand(0); - if (operand->opcode() == HloOpcode::kTuple) { - HloInstruction* forwarded_instruction = - operand->mutable_operand(instruction->tuple_index()); - VLOG(4) << "Replacing uses of " << instruction->ToString() - << " with " << forwarded_instruction->ToString(); - ABSL_RETURN_IF_ERROR( - instruction->ReplaceAllUsesWith(forwarded_instruction)); - computation_modified = true; - } - } else if (instruction->opcode() == HloOpcode::kTuple) { - // Replace Tuple(GetTupleElement(x), ..., GetTupleElement(x)) pattern - // with x. - bool can_replace = instruction->operand_count() > 0 && - instruction->operand(0)->opcode() == - HloOpcode::kGetTupleElement && - instruction->operand(0) - ->operand(0) - ->shape() - .tuple_shapes() - .size() == instruction->operand_count(); - for (int operand_number = 0; - operand_number < instruction->operand_count(); - ++operand_number) { - const HloInstruction* operand = - instruction->operand(operand_number); - if (operand->opcode() != HloOpcode::kGetTupleElement || - operand->tuple_index() != operand_number || - operand->operand(0) != instruction->operand(0)->operand(0)) { - can_replace = false; - break; - } - } - if (can_replace) { - HloInstruction* forwarded_instruction = - instruction->mutable_operand(0)->mutable_operand(0); - VLOG(4) << "Replacing uses of " << instruction->ToString() - << " with " << forwarded_instruction->ToString(); - ABSL_RETURN_IF_ERROR( - instruction->ReplaceAllUsesWith(forwarded_instruction)); - computation_modified = true; - } - } - } - } - } - - ABSL_RETURN_IF_ERROR(CleanupDeadFusionComputations(module_, removed_instructions)); - - RemoveAlternateMemoryAssignments(removed_instructions); - RemoveScopedMemoryAssignments(removed_instructions); - - return absl::OkStatus(); -} - -namespace { - -// An interface that is used to wrap asynchronous copies, asynchronous slices, -// and asynchronous slice concat operations, for use in MSA's scheduling -// algorithm (ScheduleAsynchronousCopies). -// -// Each AsyncCopy step represents 1 copy, 1 slice, or 1 concat. Each step -// has an optional start phase (e.g., to start a copy or slice), and a required -// done phase (e.g., to finish a copy or slice, or to perform a concat). -class AsyncCopyStep { - public: - struct StartPhase { - int64_t schedule_after_time; - HloInstruction* instruction; - }; - struct DonePhase { - int64_t schedule_before_time; - HloInstruction* instruction; - }; - - virtual ~AsyncCopyStep() = default; - - bool operator<(const AsyncCopyStep& rhs) const { - // AsyncCopySteps are scheduled in a multi-step algorithm: - // 1. ScheduleAsynchronousCopies - // - Build a list of needed AsyncCopySteps and sort it using this - // comparator. - // - Iterate over the sorted AsyncCopySteps and build the following data - // structures: - // * schedule_before_: A map from logical time to a list of - // instructions that should be inserted before that logical time, in - // order to perform the end of an async copy step. Such instructions - // include copy-dones. - // * schedule_after_: A map from logical time to a list of - // instructions that should be inserted after that logical time, in - // order to perform the start of an async copy step. Such - // instructions include copy-starts. - // 2. SetSchedule - This step uses the data structures built in the previous - // step to schedule the instructions in the schedule_before_ and - // schedule_after_ maps. - - // Given the way SetSchedule operates, the purpose of this comparator is to - // order the set of instructions to be scheduled after a given instruction, - // and independently order the set of instructions to be scheduled before a - // given instruction. - // * At logical time i, this comparator has the following effects: - // - Orders the set of instructions in schedule_before[i]: - // * Instructions in schedule_before_[i] are done instructions. Thus, by - // definition, all instructions in schedule_before_[i] have the same - // done time. - // * Amongst done instructions with the same done time, we first - // schedule one that starts first, since copies execute in FIFO order. - // * The next part of the comparator places evictions before prefetches, - // when the start and done times of the 2 copies are the same. - // Normally, we don't care about the ordering here, unless the 2 - // copies are jit copies. In the case of jit copies, we prefer a jit - // eviction to run and complete before a jit prefetch, so the jit - // prefetch can reuse the alternate memory freed by the jit eviction. - // In that case, this comparator ensures that we schedule the eviction - // done followed by the prefetch done. SetSchedule() then uses - // InsertInstructionAndEnsureOperandsInserted() to make the final - // ordering: evict start, evict done, prefetch start, prefetch done. - // - Ordering the set of instructions in schedule_after[i]: - // * Instructions in schedule_after_[i] are start instructions. Thus, by - // definition, all instructions in schedule_after_[i] have the same - // start time. - // * Amongst start instructions with the same start time, we first - // schedule one that ends first, since copies execute in FIFO order. - // * As mentioned above, the last part of the sort order (checking - // evictions vs prefetches) only matters for 2 jit copies, and the - // ordering of jit copy starts never reaches this criteria. - std::optional lhs_start_phase = start_phase(); - auto lhs_tuple = std::make_tuple( - done_phase().schedule_before_time, - (lhs_start_phase.has_value() ? lhs_start_phase->schedule_after_time - : done_phase().schedule_before_time), - destination_memory_space() == MemorySpace::kAlternate); - std::optional rhs_start_phase = rhs.start_phase(); - auto rhs_tuple = std::make_tuple( - rhs.done_phase().schedule_before_time, - (rhs_start_phase.has_value() ? rhs_start_phase->schedule_after_time - : rhs.done_phase().schedule_before_time), - destination_memory_space() == MemorySpace::kAlternate); - - return lhs_tuple < rhs_tuple; - } - - virtual HloPosition defining_position() const = 0; - - virtual std::optional start_phase() const = 0; - virtual void set_start_phase_schedule_after_time(int64_t schedule_after) = 0; - virtual DonePhase done_phase() const = 0; - virtual MemorySpace destination_memory_space() const = 0; - - protected: - AsyncCopyStep() = default; -}; - -class AsyncCopyStepForCopyAllocation : public AsyncCopyStep { - public: - explicit AsyncCopyStepForCopyAllocation(CopyAllocation* copy_allocation) - : AsyncCopyStep(), copy_allocation_(copy_allocation) {} - - ~AsyncCopyStepForCopyAllocation() override = default; - - HloPosition defining_position() const override { - return copy_allocation_->defining_position(); - } - - std::optional start_phase() const override { - StartPhase phase{copy_allocation_->copy_start_schedule_after(), - copy_allocation_->copy_start()}; - - return phase; - } - - void set_start_phase_schedule_after_time(int64_t schedule_after) override { - copy_allocation_->set_copy_start_schedule_after(schedule_after); - } - - DonePhase done_phase() const override { - return {copy_allocation_->copy_done_schedule_before(), - copy_allocation_->copy_done()}; - } - - MemorySpace destination_memory_space() const override { - return copy_allocation_->memory_space(); - } - - private: - CopyAllocation* copy_allocation_ = nullptr; -}; - -class AsyncCopyStepForSlice : public AsyncCopyStep { - public: - AsyncCopyStepForSlice(SlicedCopyAllocation* sliced_copy_allocation, - size_t slice_index) - : AsyncCopyStep(), - sliced_copy_allocation_(sliced_copy_allocation), - slice_index_(slice_index) {} - - ~AsyncCopyStepForSlice() override = default; - - HloPosition defining_position() const override { - return sliced_copy_allocation_->defining_position(); - } - - std::optional start_phase() const override { - const SlicedCopyAllocation::SliceDetail& slice_details = - sliced_copy_allocation_ - ->slice_details_sorted_by_start_time()[slice_index_]; - StartPhase phase{slice_details.copy_start_after_time, - slice_details.copy_start}; - - return phase; - } - - void set_start_phase_schedule_after_time(int64_t schedule_after) override { - sliced_copy_allocation_ - ->mutable_slice_details_sorted_by_start_time()[slice_index_] - .copy_start_after_time = schedule_after; - } - - DonePhase done_phase() const override { - const SlicedCopyAllocation::SliceDetail& slice_details = - sliced_copy_allocation_ - ->slice_details_sorted_by_start_time()[slice_index_]; - DonePhase phase{slice_details.copy_done_before_time, - slice_details.copy_done}; - - return phase; - } - - MemorySpace destination_memory_space() const override { - return sliced_copy_allocation_->memory_space(); - } - - private: - SlicedCopyAllocation* sliced_copy_allocation_ = nullptr; - size_t slice_index_; -}; - -class AsyncCopyStepForSliceConcat : public AsyncCopyStep { - public: - explicit AsyncCopyStepForSliceConcat( - SlicedCopyAllocation* sliced_copy_allocation) - : AsyncCopyStep(), sliced_copy_allocation_(sliced_copy_allocation) {} - - ~AsyncCopyStepForSliceConcat() override = default; - - HloPosition defining_position() const override { - return sliced_copy_allocation_->defining_position(); - } - - std::optional start_phase() const override { - return std::nullopt; - } - - void set_start_phase_schedule_after_time(int64_t schedule_after) override {} - - DonePhase done_phase() const override { - return {sliced_copy_allocation_->earliest_available_time(), - sliced_copy_allocation_->concat()}; - } - - MemorySpace destination_memory_space() const override { - return sliced_copy_allocation_->memory_space(); - } - - private: - SlicedCopyAllocation* sliced_copy_allocation_ = nullptr; -}; - -} // namespace - -void MemorySpaceAssignment::ScheduleAsynchronousCopies() { - VLOG(1) << "Scheduling asynchronous copies..."; - for (MemorySpace memory_space : - {MemorySpace::kDefault, MemorySpace::kAlternate}) { - std::vector> async_copy_steps; - for (auto& allocation : allocations_) { - if (allocation->memory_space() != memory_space) { - continue; - } - - if (allocation->is_copy_allocation()) { - auto copy_allocation = static_cast(allocation.get()); - async_copy_steps.push_back( - std::make_unique(copy_allocation)); - } else if (allocation->is_sliced_copy_allocation()) { - auto sliced_copy_allocation = - static_cast(allocation.get()); - for (int i = 0; i < sliced_copy_allocation - ->mutable_slice_details_sorted_by_start_time() - .size(); - ++i) { - async_copy_steps.push_back(std::make_unique( - sliced_copy_allocation, i)); - } - async_copy_steps.push_back( - std::make_unique( - sliced_copy_allocation)); - } - } - - absl::c_stable_sort( - async_copy_steps, - [](const std::unique_ptr& lhs, - const std::unique_ptr& rhs) { return *lhs < *rhs; }); - for (std::unique_ptr& async_copy_step : async_copy_steps) { - std::optional start_phase = - async_copy_step->start_phase(); - if (start_phase.has_value()) { - // If the copy start doesn't happen to be scheduled at the correct - // computation, delay it until the correct computation starts. - int64_t copy_start_schedule_after = start_phase->schedule_after_time; - - // Accessing flattened_instructions_ here without checking if it is - // nullptr is safe because this method is called before SimplifyGraph. - while ( - async_copy_step->defining_position().instruction->parent() != - flattened_instructions_[ - // We can't use -1 to index into flatten_instructions_. However, - // if we want to place the copy as first instruction, i.e., - // after the -1 scheduling position, its parent will be the same - // as the first instruction, i.e., the one at the 0th position. - std::max(0, copy_start_schedule_after)] - ->parent()) { - VLOG(4) << "Delaying CopyStart (" << copy_start_schedule_after - << " to " << (copy_start_schedule_after + 1) << ") for " - << start_phase->instruction->ToString() - << " because it is not in the correct computation."; - async_copy_step->set_start_phase_schedule_after_time( - ++copy_start_schedule_after); - } - start_phase = async_copy_step->start_phase(); - schedule_after_[start_phase->schedule_after_time].push_back( - start_phase->instruction); - } - - AsyncCopyStep::DonePhase done_phase = async_copy_step->done_phase(); - schedule_before_[done_phase.schedule_before_time].push_back( - done_phase.instruction); - } - } -} - -absl::Status MemorySpaceAssignment::SetSchedule() { - VLOG(1) << "Fixing schedule..."; - TF_RET_CHECK(module_->has_schedule()); - HloSchedule& schedule = module_->schedule(); - - struct ComputationStats { - HloInstructionSequence sequence; - absl::flat_hash_set inserted_instructions; - }; - - absl::flat_hash_map - computation_to_stats; - - for (const HloComputation* computation : - module_->MakeNonfusionComputations()) { - // Parallel computations aren't in the schedule and don't need to be - // modified. - if (!computations_in_schedule_.contains(computation)) { - if (computation->IsAsyncComputation()) { - VLOG(4) << "Created a dummy schedule for async computation " - << computation->name(); - schedule.GetOrCreateSequence(computation); - continue; - } - VLOG(4) << "Not scheduling " << computation->name() - << " because it's not in the schedule."; - continue; - } - TF_RET_CHECK(schedule.is_computation_scheduled(computation)); - computation_to_stats[computation] = {}; - } - - // This set contains instructions that should be ignored when iterating - // through `flattened_instructions_` to build the new schedule. MSA adds new - // asynchronous copy instructions (e.g., `copy-start`, `copy-done`) and - // decides to schedule them before or after a certain instruction as an - // anchor. For each instruction in the schedule, it first schedules the - // instructions that are supposed to be scheduled before the current - // instruction, then the current instruction and finally the instructions that - // are supposed to be scheduled after the current instruction. If the "anchor" - // instruction itself is scheduled relative to another instruction, we skip - // it when iterating over the original schedule. - absl::flat_hash_set pass_over_instructions; - for (const auto& [_, custom_call_prefetch_details] : - options_.hlo_position_to_custom_call_prefetch_details) { - for (const auto& custom_call_prefetch_detail : - custom_call_prefetch_details) { - pass_over_instructions.insert(custom_call_prefetch_detail.prefetch_start); - pass_over_instructions.insert(custom_call_prefetch_detail.prefetch_done); - for (const auto& intermediate_instruction : - custom_call_prefetch_detail.intermediate_instructions) { - pass_over_instructions.insert(intermediate_instruction); - } - } - } - - // Create the schedule for all computations at the same time. As we iterate - // through the original instructions. - // - // Iterate through the flattened schedule, doing the following: - // 1. First schedule the current instruction. - // 2. Then schedule the instructions that are supposed to be scheduled before - // the next instruction. This typically includes things like - // copy-dones. - // 3. Finally schedule the instructions that are supposed to be scheduled - // after the current instruction. This typically includes things like - // copy-starts - // - // Step 2 comes before step 3, so that we can free up space in step 2 that - // can be reused in step 3. - // - // For JIT(just-in-time) copies, the copy start will be after instruction i - // and the copy done will come before instruction i+1. Normally that would - // cause a problem because we schedule dones before starts. However, when we - // insert the done with InsertInstructionAndEnsureOperandsInserted(), it will - // end up properly scheduling the start. - // - // Schedule all instructions in their respective computations. - int64_t instructions_count = flattened_instructions_.size(); - for (int64_t instruction_index = -2; instruction_index < instructions_count; - ++instruction_index) { - // Start by inserting the instruction. - if (instruction_index >= 0) { - HloInstruction* instruction = flattened_instructions_[instruction_index]; - // Insert only if it is not deleted (SimplifyGraph sets it to nullptr if - // it was deleted) and not previously inserted. Also bitcasts and tuples - // are treated specially and only inserted as a result of operand - // dependencies. - if (instruction != nullptr && - instruction->opcode() != HloOpcode::kBitcast && - instruction->opcode() != HloOpcode::kTuple && - !pass_over_instructions.contains(instruction)) { - HloComputation* computation = instruction->parent(); - if (computation_to_stats.contains(computation)) { - ComputationStats& stats = computation_to_stats[computation]; - VLOG(4) << "inst " << instruction_index << ": " - << instruction->ToString(); - if (!stats.inserted_instructions.contains(instruction)) { - InsertInstructionAndEnsureOperandsInserted( - instruction, &stats.sequence, &stats.inserted_instructions); - } - } - } - } - - // Insert the scheduler_before_ instructions, e.g., copy dones. - auto insts_before_iter = schedule_before_.find(instruction_index + 1); - if (insts_before_iter != schedule_before_.end()) { - for (HloInstruction* new_instruction : insts_before_iter->second) { - HloComputation* computation = new_instruction->parent(); - if (computation_to_stats.contains(computation)) { - ComputationStats& stats = computation_to_stats[computation]; - VLOG(4) << "before " << instruction_index + 1 << ": " - << new_instruction->ToString(); - InsertInstructionAndEnsureOperandsInserted( - new_instruction, &stats.sequence, &stats.inserted_instructions); - } - } - } - - // Insert the scheduler_after_ instructions, e.g., copy starts. - auto insts_after_iter = schedule_after_.find(instruction_index); - if (insts_after_iter != schedule_after_.end()) { - for (HloInstruction* new_instruction : insts_after_iter->second) { - HloComputation* computation = new_instruction->parent(); - if (computation_to_stats.contains(computation)) { - VLOG(4) << "after " << instruction_index << ": " - << new_instruction->ToString(); - ComputationStats& stats = computation_to_stats[computation]; - InsertInstructionAndEnsureOperandsInserted( - new_instruction, &stats.sequence, &stats.inserted_instructions); - } - } - } - } - - for (auto& [computation, stats] : computation_to_stats) { - // Parallel computations aren't in the schedule and don't need to be - // modified. - VLOG(4) << "Scheduling: " << computation->ToString(); - - // For rare cases where the original sequence is empty, ensure the root - // instruction and its dependencies are scheduled. - InsertInstructionAndEnsureOperandsInserted(computation->root_instruction(), - &stats.sequence, - &stats.inserted_instructions); - - CHECK_EQ(stats.sequence.size(), computation->instruction_count()) - << "New sequence for computation " << computation->name() << " has " - << stats.sequence.size() << " instructions, expects " - << computation->instruction_count() << "."; - schedule.set_sequence(computation, stats.sequence); - } - - ABSL_RETURN_IF_ERROR(schedule.Update()); - - return absl::OkStatus(); -} - -namespace { - -bool IsConcatBitcastCustomCall(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kCustomCall && - instruction->custom_call_target() == - memory_space_assignment::kConcatBitcastCustomCall; -} - -// If a use is a ConcatBitcastCustomCall, we add the uses of the -// ConcatBitcastCustomCall's output buffer to the list of uses for the -// current value. Since the instructions that have ConcatBitcastCustomCall as a -// user are slices, this means we are considering all uses of the concatenated -// buffer as uses of the original slices. -std::vector GetUsesAndExtendIfConcatBitcast( - const HloValue* value, const HloDataflowAnalysis& dataflow_analysis) { - std::vector uses; - for (const HloUse& use : value->GetUses()) { - if (IsConcatBitcastCustomCall(use.instruction)) { - const HloValue& concat_bitcast_value = - dataflow_analysis.GetUniqueValueAt(use.instruction); - absl::c_copy(concat_bitcast_value.GetUses(), std::back_inserter(uses)); - } else { - uses.push_back(use); - } - } - return uses; -} - -// If a value is used by a ConcatBitcastCustomCall, we extend the time bound to -// represent the time bound of the concatenated value. -HloLiveRange::LiveRangeBounds GetTimeBoundAndExtendIfConcatBitcast( - const HloValue* value, const HloDataflowAnalysis& dataflow_analysis, - const HloLiveRange& hlo_live_range) { - HloLiveRange::LiveRangeBounds time_bound = - hlo_live_range.buffer_live_ranges().at(value); - for (const HloUse& use : value->GetUses()) { - if (IsConcatBitcastCustomCall(use.instruction)) { - const HloValue& concat_bitcast_value = - dataflow_analysis.GetUniqueValueAt(use.instruction); - const HloLiveRange::LiveRangeBounds& concat_time_bound = - hlo_live_range.buffer_live_ranges().at(&concat_bitcast_value); - time_bound.start = std::min(time_bound.start, concat_time_bound.start); - time_bound.end = std::max(time_bound.end, concat_time_bound.end); - if (hlo_live_range.instruction_schedule().at( - time_bound.end_position.instruction) < - hlo_live_range.instruction_schedule().at( - concat_time_bound.end_position.instruction)) { - time_bound.end_position = concat_time_bound.end_position; - } - } - } - return time_bound; -} - -} // namespace - -absl::Status MemorySpaceAssignment::VerifyAndExportHeapSimulatorTrace( - const HloAliasAnalysis& alias_analysis, - std::vector* alt_mem_bytes_occupied) { - VLOG(1) << "Verifying..."; - ABSL_ASSIGN_OR_RETURN(std::unique_ptr hlo_live_range, - HloLiveRange::Run(module_->schedule(), alias_analysis, - module_->entry_computation())); - - BufferIntervalTree interval_tree; - absl::flat_hash_map seen_buffers; - // The key for events is: time, is_free, value_id. This is so that the events - // are sorted first by time, then within the same time, allocations are sorted - // earlier than frees, and finally the value id as a tie breaker. - std::map, - std::tuple> - events; - - auto add_allocation_and_verify = [&](int64_t start_time, int64_t end_time, - const HeapSimulator::Chunk& chunk, - const HloValue* value) -> absl::Status { - if (IsConcatBitcastCustomCall(value->instruction())) { - return absl::OkStatus(); - } - events[std::make_tuple(start_time, /*is_free=*/false, value->id())] = - std::make_tuple(value, chunk, HeapSimulatorTrace::Event::ALLOC); - events[std::make_tuple(end_time, /*is_free=*/true, value->id())] = - std::make_tuple(value, chunk, HeapSimulatorTrace::Event::FREE); - - // Get the chunks overlapping in time and search if they overlap in space - // as well. - // TODO(berkin): For now checking against end_time - 1 (exclusive), but we - // really should check against end_time (inclusive) for cases where the - // operand can't share buffer with user (see - // HloDataflowAnalysis::CanShareOperandBufferWithUser). - for (const HeapSimulator::Chunk& overlapping_chunk : - interval_tree.ChunksOverlappingInTime(start_time, end_time - 1)) { - if (chunk.OverlapsWith(overlapping_chunk)) { - return Internal( - ("Value %s (%d, %d) off: %d size: %d overlaps with another chunk" - " off: %d size: %d"), - value->ToShortString(), start_time, end_time, chunk.offset, - chunk.size, overlapping_chunk.offset, overlapping_chunk.size); - } - } - interval_tree.Add(start_time, end_time - 1, chunk); - return absl::OkStatus(); - }; - - // Go through all instructions in the module to ensure CopyStart/CopyDone - // instructions copy between alternate memory and default memory. - for (const HloComputation* computation : - module_->MakeNonfusionComputations()) { - for (const HloInstruction* instruction : computation->instructions()) { - if (instruction->opcode() == HloOpcode::kCopyStart) { - int64_t from_memory_space = - ShapeUtil::GetSubshape(instruction->shape(), {1}) - .layout() - .memory_space(); - int64_t to_memory_space = - ShapeUtil::GetSubshape(instruction->shape(), {0}) - .layout() - .memory_space(); - CHECK_NE(from_memory_space, to_memory_space) - << "Asynchronous copy to the same memory space: " - << instruction->ToString(); - } - } - } - - for (const auto& position_and_chunk : preset_assignments_->chunks()) { - const HloPosition& position = position_and_chunk.first; - const HeapSimulator::Chunk& chunk = position_and_chunk.second; - const HloBuffer& buffer = - alias_analysis.GetUniqueBufferAt(position.instruction, position.index); - auto [it, inserted] = seen_buffers.emplace(buffer.id(), chunk); - if (!inserted) { - CHECK_EQ(it->second.offset, chunk.offset) - << "Multiple conflicting preset assignment offsets for buffer: " - << buffer.ToString() << ", pos: " << position.ToString() - << ", existing off: " << it->second.offset - << ", new off: " << chunk.offset; - CHECK_EQ(it->second.size, chunk.size) - << "Multiple conflicting preset assignment sizes for buffer: " - << buffer.ToString() << ", pos: " << position.ToString() - << ", existing size: " << it->second.size - << ", new size: " << chunk.size; - continue; - } - - for (const HloValue* value : buffer.values()) { - // When an asynchronous computation is called by multiple callers (e.g. - // in prologue, loop body, and epilogue), HloDataflowAnalysis associates - // the single async computation's root HloValue with all invocations. - // Doing so causes the HloValue's live range to span from the first - // async-start to the last async-done: - // - // Time: 0 --------- 10 --------- 28 --------- 38 --------- 49 - // Caller values: [== 10..28 ==] [== 38..49 ==] - // Async root: [================ 10..49 ===============] - // ^ (gap) ^ - // - // The actual discrete execution intervals are already accurately tracked - // and checked for the respective caller instructions (async-start, body - // parameters, async-done, etc.). Thus, skipping validation on the async - // root HloValue does not bypass validation. - if (value->instruction()->parent()->IsAsyncComputation() && - value->instruction()->parent()->caller_instructions().size() > 1) { - continue; - } - const HloLiveRange::LiveRangeBounds& time_bound = - GetTimeBoundAndExtendIfConcatBitcast( - value, alias_analysis.dataflow_analysis(), *hlo_live_range); - const HloInstruction* last_use_instruction = nullptr; - int64_t last_use_time = time_bound.start; - std::vector uses = GetUsesAndExtendIfConcatBitcast( - value, alias_analysis.dataflow_analysis()); - for (const HloUse& use : uses) { - int64_t use_time = - hlo_live_range->instruction_schedule().at(use.instruction); - if (use_time > last_use_time) { - last_use_time = use_time; - last_use_instruction = use.instruction; - } - } - - std::function - split_conditional_buffer; - split_conditional_buffer = - [&](const HloInstruction* use_instruction, int64_t start_time, - int64_t end_time, - absl::string_view indent_string) -> absl::Status { - // Special case when verifying conditional: we internally split the - // use of alternate memory in conditionals, so fish them out from the - // conditionals. - VLOG(3) << indent_string - << "Splitting conditional buffer: " << buffer.ToString() - << " value: " << value->ToShortString() << ": (" << start_time - << ", " << end_time << ") off: " << chunk.offset - << ", size: " << chunk.size; - int64_t earliest_computation_start_time = end_time; - for (const HloComputation* called_computation : - use_instruction->called_computations()) { - int64_t computation_start_time = - hlo_live_range->computation_span_times() - .at(called_computation) - .start; - earliest_computation_start_time = - std::min(earliest_computation_start_time, computation_start_time); - int64_t last_use_time = -1; - const HloInstruction* last_use_instruction = nullptr; - std::vector uses = GetUsesAndExtendIfConcatBitcast( - value, alias_analysis.dataflow_analysis()); - for (const HloUse& use : uses) { - int64_t use_time = - hlo_live_range->instruction_schedule().at(use.instruction); - if (use.instruction->parent() == called_computation && - use_time > last_use_time) { - last_use_time = use_time; - last_use_instruction = use.instruction; - } - } - if (last_use_time != -1) { - VLOG(3) << indent_string - << " computation: " << called_computation->name() << ": (" - << computation_start_time << ", " << last_use_time << ")"; - CHECK(last_use_instruction); - last_use_time = std::min(last_use_time, end_time); - if (last_use_instruction->opcode() == HloOpcode::kConditional) { - // The last use is another (nested) conditional. Call this - // function recursively. - ABSL_RETURN_IF_ERROR(split_conditional_buffer( - last_use_instruction, computation_start_time, last_use_time, - absl::StrCat(indent_string, " "))); - } else { - ABSL_RETURN_IF_ERROR(add_allocation_and_verify( - computation_start_time, last_use_time, chunk, value)); - } - } - } - VLOG(3) << indent_string << " from beginning until first computation: (" - << start_time << ", " << (earliest_computation_start_time - 1) - << ")"; - ABSL_RETURN_IF_ERROR(add_allocation_and_verify( - start_time, earliest_computation_start_time - 1, chunk, value)); - return absl::OkStatus(); - }; - - if (last_use_instruction && - last_use_instruction->opcode() == HloOpcode::kConditional) { - ABSL_RETURN_IF_ERROR(split_conditional_buffer( - last_use_instruction, time_bound.start, time_bound.end, " ")); - } else if (last_use_instruction && - last_use_instruction->opcode() == HloOpcode::kWhile && - value->instruction()->parent() != - last_use_instruction->while_body() && - value->instruction()->parent() != - last_use_instruction->while_condition() && - absl::c_none_of(uses, [&](const HloUse& use) { - return use.instruction->parent() == - last_use_instruction->while_body() || - use.instruction->parent() == - last_use_instruction->while_condition(); - })) { - // When a value defined in the caller computation is passed into a while - // loop and is not directly used inside the loop (i.e., it is not - // loop-invariant, so the loop parameter defines a new HloValue), its - // last use in the caller computation is the while instruction itself - // (which is scheduled after while_body in HloLiveRange). - // However, the caller's operand is only live in the parent computation - // until the while loop begins execution. Once the while loop starts, - // the buffer is accessed inside the while body via the body parameter, - // so we truncate the caller value's live range to end before the while - // computations start (analogous to split_conditional_buffer above). - int64_t cond_start = hlo_live_range->computation_span_times() - .at(last_use_instruction->while_condition()) - .start; - int64_t body_start = hlo_live_range->computation_span_times() - .at(last_use_instruction->while_body()) - .start; - int64_t earliest_while_computation_start_time = - std::min(cond_start, body_start); - CHECK_LT(time_bound.start, earliest_while_computation_start_time); - ABSL_RETURN_IF_ERROR(add_allocation_and_verify( - time_bound.start, earliest_while_computation_start_time - 1, chunk, - value)); - } else if (!value->GetUses().empty()) { - last_use_time = std::min(last_use_time, time_bound.end); - VLOG(3) << " buffer: " << buffer.ToString() - << " value: " << value->ToShortString() << ": (" - << time_bound.start << ", " << last_use_time - << ") off: " << chunk.offset << ", size: " << chunk.size; - ABSL_RETURN_IF_ERROR(add_allocation_and_verify(time_bound.start, - last_use_time, chunk, value)); - } - } - } - - HeapSimulatorTrace* heap_trace = - &preset_assignments_ - ->assignment_information_for_space(options_.alternate_memory_space) - ->heap_simulator_trace; - int64_t memory_usage = 0; - int64_t max_memory_usage = 0; - int64_t prev_time = 0; - int64_t prev_memory_usage = 0; - if (alt_mem_bytes_occupied) { - // Populate alt_mem_bytes_occupied with -1, for each instruction. - alt_mem_bytes_occupied->resize( - hlo_live_range->flattened_instruction_sequence().instructions().size(), - -1); - } - for (const auto& event : events) { - int64_t time; - bool is_free; - int64_t buffer_id; - std::tie(time, is_free, buffer_id) = event.first; - const HloValue* value; - HeapSimulator::Chunk chunk; - HeapSimulatorTrace::Event::Kind kind; - std::tie(value, chunk, kind) = event.second; - HeapSimulatorTrace::Event* heap_trace_event = heap_trace->add_events(); - heap_trace_event->set_kind(kind); - heap_trace_event->set_buffer_id(buffer_id); - // When building heap simulation trace events, unwrap asynchronous - // instructions so that profile visualizations display the meaningful - // underlying operation name. - const HloInstruction* trace_instr = value->instruction(); - if (trace_instr->IsAsynchronous() && - trace_instr->async_wrapped_instruction() != nullptr) { - trace_instr = trace_instr->async_wrapped_instruction(); - } - *heap_trace_event->mutable_instruction_name() = - std::string(trace_instr->name()); - *heap_trace_event->mutable_computation_name() = - std::string(value->instruction()->parent()->name()); - - if (prev_time != time) { - VLOG(2) << "Memory usage: " << std::max(memory_usage, prev_memory_usage) - << " at time: " << prev_time << " (" - << hlo_live_range->flattened_instruction_sequence() - .instructions() - .at(prev_time) - ->name() - << ")"; - prev_time = time; - prev_memory_usage = memory_usage; - } - if (kind == HeapSimulatorTrace::Event::ALLOC) { - memory_usage += chunk.size; - } else { - CHECK_EQ(kind, HeapSimulatorTrace::Event::FREE); - memory_usage -= chunk.size; - } - prev_memory_usage = std::max(prev_memory_usage, memory_usage); - max_memory_usage = std::max(max_memory_usage, memory_usage); - if (alt_mem_bytes_occupied) { - // Update known alt mem usage. - (*alt_mem_bytes_occupied)[time] = memory_usage; - } - } - if (alt_mem_bytes_occupied) { - // Replace the -1s in alt_mem_bytes_occupied with the previous alt memory - // usage. - int64_t prev_bytes = 0; - for (int64_t i = 0; i < alt_mem_bytes_occupied->size(); ++i) { - if ((*alt_mem_bytes_occupied)[i] == -1) { - (*alt_mem_bytes_occupied)[i] = prev_bytes; - } - prev_bytes = (*alt_mem_bytes_occupied)[i]; - } - } - - VLOG(1) << "Max memory usage ignoring fragmentation: " << max_memory_usage; - - return absl::OkStatus(); -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h index 2de772dca1499b..e26d7b83cd8495 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h @@ -195,25 +195,11 @@ Useful logging and error messages #include #include -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" #include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/layout.h" #include "xla/service/heap_simulator/heap_simulator.h" #include "xla/service/hlo.pb.h" #include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/util.h" namespace xla { namespace memory_space_assignment { @@ -306,163 +292,6 @@ class PresetAssignments { std::string instruction_schedule_str_; }; -// MemorySpaceAssignment assigns memory spaces (default or alternate) to each -// instruction in the module. It will greedily try placing as as many values in -// the alternate memory space as possible. It uses the heap simulator to -// determine the actual allocation offsets of values in the alternate memory -// space to account for fragmentation. The default memory space is assumed to be -// large enough to hold the values that could not be placed in the alternate -// memory space. -class MemorySpaceAssignment { - public: - // Statistics of asynchronous copies. - struct AsyncCopyStats { - // Includes both async copies and async sliced copies. - int64_t max_outstanding_async_copies = 0; - // Includes both async copies and async sliced copies. - int64_t num_prefetches = 0; - int64_t num_sliced_prefetches = 0; - int64_t num_sliced_prefetch_slices = 0; - int64_t prefetch_bytes = 0; - int64_t num_evictions = 0; - int64_t eviction_bytes = 0; - }; - - virtual ~MemorySpaceAssignment() = default; - - // Runs the MemorySpaceAssignment pass. - static absl::StatusOr> Run( - HloModule* module, const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis, const AliasInfo* alias_info, - const Options& options); - - // Calculates asynchronous copy statistics. - absl::StatusOr CalculateAsyncCopyStats( - const HloDataflowAnalysis& dataflow_analysis) const; - - // Verify that allocations_ are free of overlapping Allocations in time and - // space. This is a post-processing step called after all allocations have - // been finalized, before the async copies get scheduled. - absl::Status VerifyAllocations() const; - - // Verify that the memory space assignment is free of overlapping buffers and - // export heap simulator trace to be used by buffer_assignment. - // - // If alt_mem_bytes_occupied is not null, it will be populated with the number - // of bytes occupied in the alternate memory space at each instruction time. - absl::Status VerifyAndExportHeapSimulatorTrace( - const HloAliasAnalysis& alias_analysis, - std::vector* alt_mem_bytes_occupied = nullptr); - - static constexpr absl::string_view kName = "memory-space-assignment"; - - protected: - // Main driver of the memory space assignment pass. - virtual absl::StatusOr> - RunMemorySpaceAssignment(const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis); - - // Finds an AllocationSequence for placing buffers in alternate memory using - // the MsaAlgorithm algorithm. Must be set before Process() is called. - virtual absl::Status FindAllocationSequence( - const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis); - - const Options& options() const { return options_; } - - MemorySpaceAssignment(HloModule* module, const AliasInfo* alias_info, - const Options& options, - const HloLiveRange& hlo_live_range) - : module_(module), - alias_info_(alias_info), - options_(options), - flattened_instructions_(hlo_live_range.flattened_instruction_sequence() - .instructions() - .begin(), - hlo_live_range.flattened_instruction_sequence() - .instructions() - .end()), - computations_in_schedule_(), - preset_assignments_(std::make_unique()) { - for (const auto& computation_and_bound : - hlo_live_range.computation_span_times()) { - computations_in_schedule_.insert(computation_and_bound.first); - } - } - - AllocationSequence allocations_; - - HloModule* module() { return module_; } - - private: - // A struct that represents the source of scoped alternate memory. It can be - // either for an instruction or for post-module operations. - struct ScopedMemorySource { - static ScopedMemorySource ForInstruction(HloInstruction* instruction); - static ScopedMemorySource ForPostModule(); - - std::string ToString() const; - - bool is_post_module = false; - HloInstruction* instruction = nullptr; - }; - - // Process calls Process methods of the allocations after the allocations have - // been finalized. - absl::Status Process(const HloLiveRange& hlo_live_range, - const HloAliasAnalysis& alias_analysis); - - // Process() might have altered the computation graph by inserting kTuple and - // kGetTupleElement instructions. SimplifyGraph performs a simple DCE and - // tuple simplification operation (e.g., given GetTupleElement(Tuple(a, b), - // 1), simply forwards b). Runs to fixed point. - absl::Status SimplifyGraph(); - - // Places copy-start and copy-done instructions in the schedule, according to - // the schedule_before_ and schedule_after_ data structures, making necessary - // adjustments to the schedule to ensure correctness. - absl::Status SetSchedule(); - - // Export the alternate memory assignments to the PresetAssignments and color - // the HLO graph with the determined memory spaces. - absl::Status ExportAndColorBuffers(const HloAliasAnalysis& alias_analysis); - - // Schedules asynchronous copies and ensures that the CopyStarts and their - // corresponding CopyDones follow the same order. - void ScheduleAsynchronousCopies(); - - // Remove the positions and chunks associated with instructions, from - // alternate_memory_assignments_. - void RemoveAlternateMemoryAssignments( - const absl::flat_hash_set& instructions); - - // Remove the positions and chunks associated with instructions, from - // scoped_memory_assignments_. - void RemoveScopedMemoryAssignments( - const absl::flat_hash_set& instructions); - - HloModule* module_; - // Backend specific aliasing information. - const AliasInfo* alias_info_; - const Options& options_; - std::vector flattened_instructions_; - absl::flat_hash_set computations_in_schedule_; - std::unique_ptr preset_assignments_; - std::vector> - alternate_memory_assignments_; - // Maps from a defining position to a shape if the tensor is split. - absl::flat_hash_map split_map_; - std::vector> - scoped_memory_assignments_; - int64_t alternate_memory_size_ = 0; - - // These maps hold vectors of new instructions that need to be scheduled after - // (or before) the instruction index in the key. FixSchedule uses these maps - // to modify and fix the schedule. - absl::flat_hash_map> schedule_after_; - absl::flat_hash_map> schedule_before_; -}; - } // namespace memory_space_assignment } // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto deleted file mode 100644 index db0f1949eacfe9..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto +++ /dev/null @@ -1,261 +0,0 @@ -/* Copyright 2023 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -syntax = "proto3"; - -package xla.memory_space_assignment; - -// Memory space assignment options for slicing prefetches into smaller -// asynchronous copies, reducing prefetch memory allocation pressure. -// -// No prefetch slicing is performed if max_slices == 0. -// -// TODO(b/275905276): Consider adding another option that indicates that we want -// slices of a certain size, rather than just always creating max_slices. -message SlicedPrefetchOptions { - // The maximum number of slices into which to slice a prefetch. - uint32 max_slices = 1; - - // The minimum tensor size in bytes that we will attempt to slice. - uint64 min_bytes = 2; - - // This option should never be set to true in production. When this is true, - // we will crash if we propose a slice (other than the final slice) with a - // size that is not a multiple of the required hardware alignment. Otherwise, - // we will choose not to slice such situations, which is always safe. - bool fail_on_non_alignment_boundary_slice_proposal = 3; - - // The threshold for max_slices after which we limit the permutations of slice - // times that we try when placing a sliced allocation. - uint32 all_slice_time_permutations_threshold = 4; - - // The preferred slize size for MSA sliced prefetches. 0 means there is no - // preferred slice size, in which case, we'll try to slice into max_slices. - uint64 preferred_slice_size = 5; -} - -// Options for memory-bound loop optimizations in memory space assignment. If -// enabled, this pass can optimize memory-bound unrolled loops to maximize the -// bandwidth utilized and minimize the execution time. -message MemoryBoundLoopOptimizerOptions { - // Enable the memory-bound loop optimizations. - optional bool enabled = 1; - - // The desired ratio of overlapped operations that is sufficient to overlap - // prefetches with. If this value is 1, the algorithm will try to fully - // overlap the prefetches with other compute, if less than 1, the algorithm - // may schedule prefetches such that some of the prefetch is not overlapped, - // so may become critical. For example, if this value is 0.5, we are willing - // for the prefetch time to take up to 2X of the overlapped computation time. - optional float desired_copy_ratio = 2; - - // If true, the algorithm allows a fully pipelined prefetch to be scheduled - // even if the copy resources haven't reached the desired copy ratio. A fully - // pipelined prefetch starts the same time as its counterpart in the previous - // iteration finishes. - optional bool allow_unsatisfied_fully_pipelined_prefetch = 3; - - // The minimum number of iterations that the loop needs to be unrolled for the - // memory-bound loop optimizer to kick in. - optional float min_num_iterations = 4; -} - -message TupleShapeIndex { - repeated int64 index = 1; -} - -// A message to filter operands in an HLO schedule, that can be used to override -// compiler behaviour like altering schedule etc. -message HloOperandFilter { - // Regex to match instruction name. - optional string instruction_name_regex = 1; - // Set if filtering operands of an instruction. - optional int64 operand_number = 2; - // If filtering operands based on size in bytes. - optional int64 size_gte = 3; - // If filtering operands based on size in bytes. - optional int64 size_lte = 4; - // If operand of an instruction is a tuple and indexing into the tuple is - // required. - optional TupleShapeIndex tuple_index = 5; - // Regex to match the entire instruction HLO. The HLO string is constructed - // using default HloPrintOptions. Refer to the HloPrintOptions class in - // hlo_instruction.h to know more about the format of the HLO string used for - // matching. - optional string instruction_regex = 6; -} - -// Options to override preferred prefetch time for an operand. -message PreferredPrefetchOverrideOptions { - reserved "after_instruction_name", // 2 - "before_instruction_name"; // 3 - reserved 2, 3; - oneof options { - // A value X in [0, 1] that tells us the preferred prefetch time is the - // fraction X through the live range. For example, .5 will set the - // preferred prefetch time to the middle of live range. - float prefetch_eagerness = 1; - // Preferred prefetch time is set to after the first instruction that - // matches the filter. - HloPositionMatcher after_instruction = 4; - // Preferred prefetch time is set to before the first instruction that - // matches the filter. - HloPositionMatcher before_instruction = 5; - // Preferred prefetch time set directly to a logical schedule time. - int64 logical_time = 6; - } - // If true, the picker and allocator must ONLY attempt this exact time (no - // outward search). If it cannot be scheduled at this exact time, fail - // compilation. - optional bool strict_timing = 7; -} - -// Filters operands in an HLO schedule and overrides preferred prefetch times -// for those operands according to an override strategy specified in -// override_options. -message PreferredPrefetchOverride { - optional HloOperandFilter hlo_operand_filter = 1; - optional xla.memory_space_assignment.PreferredPrefetchOverrideOptions - override_options = 2; - // If true, any failure to fulfill the action results in a fatal error. - optional bool fail_on_unsatisfied_override = 3; -} - -// Encloses chained override configs. The first config has highest precedence -// and so on. -message PreferredPrefetchOverrides { - repeated PreferredPrefetchOverride overrides = 1; - // Global flag to fail compilation if any override is unsatisfied. - optional bool fail_on_unsatisfied_override = 2; -} - -// Specifies details on how to randomally select HloPositions for perturbation. -// -// Each HloPosition is randomally assigned (based on the seed) to a value in the -// range [0...1). If it is within the range -// [selection_range_begin...selection_range_end) it will have its prefetch -// priority overridden based on the value of MsaSortOrderOverrideOptions. -// Anything outside this range will be left as-is. -// -// We use the selection range (rather than just supplying the probability) to be -// able to binary search in the case where there's just one instruction that -// needs to have its priority fixed. -// -// The probability of any HloPosition to be perturbed is -// selection_range_end-selection_range_begin. -// -// NOTE: different HloPositions with the same computation graph (the same -// instruction fingerprint) will be considered as identical, and hence will -// receive the same priority from the HloRandomFilter. -message HloRandomFilter { - // Defaults to 1234 - optional int64 seed = 1; - // Defaults to 0. - optional double selection_range_begin = 2; - // Defaults to 1. - optional double selection_range_end = 3; -} - -// A message that identifies one or more HloPositions. -message HloPositionMatcher { - // Regex to match the entire instruction HLO. The HLO string is constructed - // using default HloPrintOptions. Refer to the HloPrintOptions class in - // hlo_instruction.h to know more about the format of the HLO string used for - // matching. - optional string instruction_regex = 1; - // Regex to match instruction name. - optional string instruction_name_regex = 2; - // If output of an instruction is a tuple and indexing into the - // tuple is required. - optional TupleShapeIndex tuple_index = 3; - // Filters instructions with output size in bytes greater or equal to a value. - optional int64 size_gte = 4; - // Filters instructions with output size in bytes less or equal to a value. - optional int64 size_lte = 5; - // Filters instructions that have a use that matches the filter. - optional HloOperandFilter hlo_use_filter = 6; - // Filters instructions randomly. - optional HloRandomFilter hlo_random_filter = 7; -} - -// Options to override preferred prefetch time for an operand. -message MsaSortOrderOverrideOptions { - oneof options { - // Assign alternate memory to the filtered buffer before other buffers. If - // multiple buffers are to be assigned first (within the same override - // config) other tie breakers and stable sort order will take effect. - bool assign_first = 1; - // Assign alternate memory to the filtered buffer after other buffers. If - // multiple buffers are to be assigned last (within the same override - // config) other tie breakers and stable sort order will take effect. - bool assign_last = 2; - // Set the specific priority value, where lower values means "before" higher - // values. If multiple buffers are to be assigned the same value, other tie - // breakers and stable sort order will take effect. - int64 assign_value = 3; - } -} - -// Specifies details on how to override the sort order for matching -// HloPositions. -message MsaSortOrderOverride { - optional HloPositionMatcher hlo_position_matcher = 1; - optional xla.memory_space_assignment.MsaSortOrderOverrideOptions - override_options = 2; - optional bool apply_to_cross_program_prefetches = 3; -} - -// Encloses chained override configs. The first config has highest precedence -// and so on. -message MsaSortOrderOverrides { - repeated MsaSortOrderOverride overrides = 1; -} - -// Action to pin a buffer in alternate memory across its entire lifetime. -message PinInAlternateMemoryAction {} - -// Action to keep a buffer or operand use strictly in default memory (HBM). -message KeepInDefaultMemoryAction {} - -// Unified override rule for tensor and operand placement and timing. Overrides -// specified via MsaTensorOverride are always strictly enforced; compilation -// fails if an override cannot be fulfilled. -message MsaTensorOverride { - // Target filter for operand use or buffer position. - optional HloOperandFilter hlo_operand_filter = 1; - optional HloPositionMatcher hlo_position_matcher = 2; - - oneof action { - PreferredPrefetchOverrideOptions prefetch = 3; - PinInAlternateMemoryAction pin_in_alternate_memory = 4; - KeepInDefaultMemoryAction keep_in_default_memory = 5; - } -} - -// Encloses chained tensor override configs. -message MsaTensorOverrides { - repeated MsaTensorOverride overrides = 1; -} - -// Expanded scoped alternate memory is a feature used at the end of MSA, in -// in which we attempt to expand the size of allocated scoped alternate memory -// buffers to the largest contiguous open space available. -message ExpandedScopedAlternateMemoryMode { - enum Value { - UNDEFINED = 0; - DISABLED = 1; - ENABLED = 2; - } -} diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc deleted file mode 100644 index eef1a44762504f..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc +++ /dev/null @@ -1,20049 +0,0 @@ -/* Copyright 2019 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/memory_space_assignment.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/status/status_matchers.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" -#include "absl/strings/str_replace.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/comparison_util.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_module.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/ir/hlo_print_options.h" -#include "xla/hlo/ir/hlo_schedule.h" -#include "xla/hlo/testlib/verified_hlo_module.h" -#include "xla/hlo/transforms/simplifiers/instruction_hoister.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/hlo/utils/hlo_matchers.h" -#include "xla/layout.h" -#include "xla/layout_util.h" -#include "xla/literal_util.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/heap_simulator/allocation_block.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo.pb.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/algorithm.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/allocation_value.h" -#include "xla/service/memory_space_assignment/best_fit_repacker.h" -#include "xla/service/memory_space_assignment/buffer_interval_comparator.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/memory_space_assignment_test_base.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/service/memory_space_assignment/prefetch_interval_picker.h" -#include "xla/service/memory_space_assignment/repacking.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/service/memory_space_assignment/testing_utils.h" -#include "xla/service/memory_space_assignment/utils.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/tests/test_utils.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/logging.h" -#include "xla/util.h" -#include "xla/xla_data.pb.h" -#include "tsl/platform/protobuf.h" // IWYU pragma: keep - -namespace xla { -namespace memory_space_assignment { -namespace { - -namespace op = xla::testing::opcode_matchers; -using Chunk = HeapSimulator::Chunk; -using ::testing::_; -using ::testing::Return; -using ::testing::UnorderedElementsAre; - -constexpr float kBytesPerSecond = 100; - -const auto& ShapeSize = HloCostAnalysis::DefaultShapeSize; - -using MemorySpaceAssignmentTest = MemorySpaceAssignmentTestBase; - -// A mock MemorySpaceAssignmentRepacker class that accepts a map of -// (start_time,offset) -> new_offset values. Using this map, the repacker -// repacks the allocations to the new_offset. -class FakeMemorySpaceAssignmentRepacker : public MemorySpaceAssignmentRepacker { - public: - explicit FakeMemorySpaceAssignmentRepacker( - absl::flat_hash_map, int64_t>& repack_map, - std::function)> check_fun = nullptr, - bool always_return_modified = false) - : MemorySpaceAssignmentRepacker(/*max_size=*/128, /*alignment=*/8), - repack_map_(repack_map), - check_fun_(check_fun), - always_return_modified_(always_return_modified) {} - - absl::StatusOr Repack( - absl::Span allocations) override { - bool modified = false; - for (AllocationBlock* block : allocations) { - absl::flat_hash_set colocations; - std::string colocations_str; - for (const AllocationBlock* colocation : block->GetColocations()) { - absl::StrAppend(&colocations_str, colocation->id, ", "); - colocations.insert(colocation->id); - } - VLOG(1) << "Alloc id: " << block->id << " time: [" - << block->inclusive_start_time << ", " << block->end_time - << "] size: " << block->size - << " init offset: " << block->initial_offset << " colocations: {" - << colocations_str << "}"; - auto it = repack_map_.find( - {block->inclusive_start_time, block->initial_offset}); - if (it != repack_map_.end()) { - modified = true; - block->offset = it->second; - } else { - block->offset = block->initial_offset; - } - for (AllocationBlock* colocation : block->GetColocations()) { - if (it != repack_map_.end()) { - colocation->offset = it->second; - } else { - colocation->offset = colocation->initial_offset; - } - } - } - if (check_fun_) { - check_fun_(allocations); - } - - return always_return_modified_ || modified; - } - - private: - // A map from (start_time, offset) to new_offset. - absl::flat_hash_map, int64_t> repack_map_; - std::function)> check_fun_; - bool always_return_modified_; -}; - -TEST_F(MemorySpaceAssignmentTest, ParameterOnly) { - // A module consisting of a single parameter. Inputs/outputs are currently - // excluded from memory space assignment. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); -} - -TEST_F(MemorySpaceAssignmentTest, Simple) { - // A simple module with a few simple instructions. Expect this to be - // transformed with CopyStart and CopyDone instructions inserted after inputs - // and before outputs. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, p1)); - HloInstruction* sub = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* mul = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, add, sub)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, add, sub, mul}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.post_module_scoped_alternate_memory_size_in_bytes = 10; - auto preset_assignments = AssignMemorySpace(module.get(), std::move(options)); - - // Check that the post-module scoped alternate memory chunk is set correctly. - ASSERT_TRUE(preset_assignments->post_module_scoped_alternate_memory_chunk() - .has_value()); - EXPECT_EQ( - preset_assignments->post_module_scoped_alternate_memory_chunk()->size, - 10); - - // Inputs and outputs are currently placed in the default memory. Everything - // else should be in the alternate memory. - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - EXPECT_THAT(mul, op::ShapeWithLayout(shape)); - EXPECT_THAT(add, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(sub, op::ShapeWithLayout(shape_in_alternate_mem)); - - // Make sure the preset assignments is sane. - EXPECT_EQ(preset_assignments->chunks().size(), 3); - EXPECT_EQ(preset_assignments->assignment_informations().size(), 1); - // Ensure the offset assigned to add and sub are different. - EXPECT_NE(preset_assignments->chunks()[0].second.offset, - preset_assignments->chunks()[1].second.offset); -} - -TEST_F(MemorySpaceAssignmentTest, BasicSplit) { - constexpr int64_t kSplitDimension = 1; - constexpr int64_t kSplitIndex = 256; - SplitConfig split_config(/*dimension=*/kSplitDimension, {kSplitIndex}); - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {256, 512}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, p1)); - HloInstruction* sub = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* mul = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, add, sub)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, add, sub, mul}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 256 * 512 * 16; - options.init_split_tree_fn = - [](const HloInstruction* instruction, - absl::flat_hash_map>* - split_map) -> absl::flat_hash_map { - if (split_map != nullptr) { - auto result = split_map->find(instruction); - if (result != split_map->end()) { - return result->second; - } - } - return {}; - }; - - options.determine_split_dimension_fn = - [&split_config]( - const HloValue& hlo_value, - absl::flat_hash_map>* - split_map) -> std::optional { - if (hlo_value.instruction()->opcode() == HloOpcode::kAdd || - hlo_value.instruction()->opcode() == HloOpcode::kSubtract) { - return split_config; - } - return std::nullopt; - }; - - auto preset_assignments = AssignMemorySpace(module.get(), std::move(options)); - - // Inputs and outputs are currently placed in the default memory. Everything - // else should be in the alternate memory. - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {256, 512}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - shape_in_alternate_mem.mutable_layout()->add_split_configs(split_config); - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - EXPECT_THAT(mul, op::ShapeWithLayout(shape)); - EXPECT_THAT(add, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(sub, op::ShapeWithLayout(shape_in_alternate_mem)); - - // Make sure the preset assignments is sane. - EXPECT_EQ(preset_assignments->chunks().size(), 3); - EXPECT_EQ(preset_assignments->assignment_informations().size(), 1); - // Ensure the offset assigned to add and sub are different. - EXPECT_NE(preset_assignments->chunks()[0].second.offset, - preset_assignments->chunks()[1].second.offset); -} - -TEST_F(MemorySpaceAssignmentTest, NegateChain) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[2], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, PinnedDefaultMemorySpace) { - absl::string_view hlo_string = R"hlo( - HloModule NegateChain, is_scheduled=true, entry_computation_layout={(f32[2,3]{1,0}, f32[2,3]{1,0:S(2)})->f32[2,3]{1,0}} - - ENTRY %NegateChain (p0: f32[2,3], p1: f32[2,3]) -> f32[2,3] { - %p0 = f32[2,3]{1,0} parameter(0) - %p1 = f32[2,3]{1,0:S(2)} parameter(1) - %negate = f32[2,3]{1,0:S(2)} negate(f32[2,3]{1,0} %p0) - %negate.1 = f32[2,3]{1,0:S(2)} negate(f32[2,3]{1,0:S(2)} %negate) - %negate.2 = f32[2,3]{1,0:S(2)} negate(f32[2,3]{1,0:S(2)} %negate.1) - %negate.3 = f32[2,3]{1,0} negate(f32[2,3]{1,0:S(2)} %negate.2) - %negate.4 = f32[2,3]{1,0:S(2)} negate(f32[2,3]{1,0} %negate.3) - %negate.5 = f32[2,3]{1,0:S(2)} negate(f32[2,3]{1,0:S(2)} %negate.4) - %negate.6 = f32[2,3]{1,0:S(2)} negate(f32[2,3]{1,0:S(2)} %negate.5) - ROOT %add = f32[2,3]{1,0} add(f32[2,3]{1,0:S(2)} %negate.6, f32[2,3]{1,0:S(2)} %p1) - })hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - XLA_VLOG_LINES(1, module->ToString()); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* p1 = FindInstruction(module.get(), "p1"); - HloInstruction* negate = FindInstruction(module.get(), "negate"); - HloInstruction* negate_1 = FindInstruction(module.get(), "negate.1"); - HloInstruction* negate_2 = FindInstruction(module.get(), "negate.2"); - HloInstruction* negate_3 = FindInstruction(module.get(), "negate.3"); - HloInstruction* negate_4 = FindInstruction(module.get(), "negate.4"); - HloInstruction* negate_5 = FindInstruction(module.get(), "negate.5"); - HloInstruction* negate_6 = FindInstruction(module.get(), "negate.6"); - HloInstruction* add = FindInstruction(module.get(), "add"); - std::vector pinned_hbm_instructions = { - p1, negate, negate_1, negate_2, negate_4, negate_5, negate_6}; - for (const HloInstruction* instruction : pinned_hbm_instructions) { - EXPECT_EQ(instruction->shape().layout().memory_space(), - kPinnedDefaultMemorySpace); - } - // Check p0 and add are in the default memory space. - EXPECT_EQ(p0->shape().layout().memory_space(), kDefaultMemorySpace); - EXPECT_EQ(add->shape().layout().memory_space(), kDefaultMemorySpace); - // Check negate_3 is in pinned to alternate memory space. - EXPECT_EQ(negate_3->shape().layout().memory_space(), kAlternateMemorySpace); - // Check that p1 is only used once at the add instruction. ie, the there is no - // copy/prefetch. - CHECK_EQ(p1->users().size(), 1); - EXPECT_EQ(p1->users()[0], add); -} - -// A simple case where the synchronous copy is actually redundant, because its -// operand ends up getting prefetched and its output is only used once, so we -// remove the sync copy. -TEST_F(MemorySpaceAssignmentTest, - SyncCopyReplacementRedundantCopyAfterPrefetch) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = f32[2,3]{1,0} negate(p1) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(negate2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[2,3]{1,0} negate(negate5) - negate7 = f32[2,3]{1,0} negate(negate6) - p0_copy = f32[2,3]{1,0} copy(p0) - ROOT add0 = f32[2,3]{1,0} add(p0_copy, negate7) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_sync_copy_replacement = true; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* add0 = FindInstruction(module.get(), "add0"); - ASSERT_NE(add0, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT(add0->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); -} - -TEST_F(MemorySpaceAssignmentTest, SyncCopyReplacementWithControlPredecessor) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = f32[2,3]{1,0} negate(p1) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(negate2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[2,3]{1,0} negate(negate5) - negate7 = f32[2,3]{1,0} negate(negate6) - p0_copy = f32[2,3]{1,0} copy(p0) - negate8 = f32[2,3]{1,0} negate(negate7) - negate9 = f32[2,3]{1,0} negate(negate8), control-predecessors={p0_copy} - ROOT add0 = f32[2,3]{1,0} add(p0_copy, negate9) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_sync_copy_replacement = true; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* add0 = FindInstruction(module.get(), "add0"); - ASSERT_NE(add0, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT(add0->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - HloInstruction* negate9 = FindInstruction(module.get(), "negate9"); - ASSERT_NE(negate9, nullptr); - const HloInstruction* copy_done = add0->operand(0); - const HloInstructionSequence& sequence = - module->schedule().sequence(module->entry_computation()); - auto find_index = [&](const HloInstruction* instruction) { - return std::distance(sequence.instructions().begin(), - absl::c_find(sequence.instructions(), instruction)); - }; - int64_t copy_done_time = find_index(copy_done); - int64_t negate9_time = find_index(negate9); - // The negate9 instruction should be scheduled after the copy done, because of - // the control dependency constraint. - EXPECT_GT(negate9_time, copy_done_time); -} - -// This is a case where we p0_copy uses and p0 uses after copy(p0) are not -// allowed to use the same async CopyAllocation. While p0 can be prefetched at -// p0_copy, but we may clobber the data if we use the same async copy used for -// prefetching to replace the sync copy p0_copy. The pattern here is that the -// sync copy operand (p0) shows up at more than one position. -TEST_F(MemorySpaceAssignmentTest, - SyncCopyReplacementWouldNeedMoreThanOneAsyncCopy) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = f32[2,3]{1,0} negate(p1) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(negate2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[2,3]{1,0} negate(negate5) - negate7 = f32[2,3]{1,0} negate(negate6) - p0_copy0 = f32[2,3]{1,0} copy(p0) - p0_copy1 = f32[2,3]{1,0} copy(p0) - negate8 = f32[2,3]{1,0} negate(negate7) - negate9 = f32[2,3]{1,0} negate(negate8) - negate10 = f32[2,3]{1,0} negate(negate9) - negate11 = f32[2,3]{1,0} negate(negate10) - negate12 = f32[2,3]{1,0} negate(negate11) - constant.1 = f32[] constant(0) - broadcast = f32[2,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - dynamic-update-slice.0 = f32[2,3] dynamic-update-slice(p0_copy0, broadcast, constant.3, constant.3) - dynamic-update-slice.1 = f32[2,3] dynamic-update-slice(p0_copy1, broadcast, constant.3, constant.3) - ROOT tuple0 = tuple(negate12, dynamic-update-slice.0, dynamic-update-slice.1) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_sync_copy_replacement = true; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - - HloInstruction* dynamic_update_slice_0 = - FindInstruction(module.get(), "dynamic-update-slice.0"); - HloInstruction* dynamic_update_slice_1 = - FindInstruction(module.get(), "dynamic-update-slice.1"); - const HloInstruction* p0_copy0_replacement = - dynamic_update_slice_0->operand(0); - const HloInstruction* p0_copy1_replacement = - dynamic_update_slice_1->operand(0); - EXPECT_THAT(p0_copy0_replacement, - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - EXPECT_THAT(p0_copy1_replacement, - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - ASSERT_NE(p0_copy0_replacement, p0_copy1_replacement); -} - -// All uses of the sync copy operand that are scheduled before the replaced sync -// copy share the allocation in alternate memory (if any). -TEST_F(MemorySpaceAssignmentTest, SyncCopyReplacementOperandHasMultipleUses) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - p0_negate0 = negate(p0) - p0_negate1 = negate(p0) - negate6 = negate(negate5) - negate7 = negate(negate6) - p0_copy = f32[2,3]{1,0} copy(p0) - add0 = add(p0_copy, p0_negate0) - ROOT tuple = tuple(negate7, add0, p0_negate1) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_sync_copy_replacement = true; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* add0 = FindInstruction(module.get(), "add0"); - HloInstruction* p0_negate0 = FindInstruction(module.get(), "p0_negate0"); - HloInstruction* p0_negate1 = FindInstruction(module.get(), "p0_negate1"); - - EXPECT_THAT(add0->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - EXPECT_THAT(p0_negate0->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - EXPECT_THAT(p0_negate1->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - ASSERT_EQ(p0_negate0->operand(0), p0_negate1->operand(0)); - ASSERT_NE(add0->operand(0), p0_negate0->operand(0)); -} - -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementAfterPrefetch) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - slice = f32[1,2,3] slice(p0), slice={[0:1], [0:2], [0:3]} - concat = f32[11,2,3] concatenate(negate7, slice), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT(concat->operand(1), op::AsyncDone(op::AsyncStart(p0))); -} - -TEST_F(MemorySpaceAssignmentTest, ViewExtendedUseTimeWalksTransitiveReaders) { - // A view (dus_view_color colored custom call) aliases its operand's storage - // without owning any, and the read side rewrite reattaches a view colored - // bitcast ON the view, so the real consumer reads the source THROUGH the - // bitcast. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[4,4]{1,0} parameter(0) - p1 = f32[100]{0} parameter(1) - p3 = f32[16]{0} parameter(2) - source = f32[4,4]{1,0} negate(p0) - view = f32[4,4]{1,0:S(5)} custom-call(source), custom_call_target="tpu_get_view" - viewbc = f32[16]{0:S(5)} bitcast(view) - negate0 = f32[100]{0} negate(p1) - negate1 = f32[100]{0} negate(negate0) - negate2 = f32[100]{0} negate(negate1) - negate3 = f32[100]{0} negate(negate2) - consumer = f32[16]{0} add(viewbc, p3) - ROOT tuple = (f32[16]{0}, f32[100]{0}) tuple(consumer, negate3) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - absl::flat_hash_map schedule; - int64_t time = 0; - for (const HloInstruction* instruction : - module->schedule() - .sequence(module->entry_computation()) - .instructions()) { - schedule[instruction] = time++; - } - const HloInstruction* view = FindInstruction(module.get(), "view"); - const HloInstruction* viewbc = FindInstruction(module.get(), "viewbc"); - const HloInstruction* consumer = FindInstruction(module.get(), "consumer"); - ASSERT_NE(view, nullptr); - ASSERT_NE(viewbc, nullptr); - ASSERT_NE(consumer, nullptr); - // The transitive walk must reach `consumer` through the view colored - // bitcast, well past the bitcast's own (early) schedule position. - EXPECT_GT(schedule.at(consumer), schedule.at(viewbc) + 1); - EXPECT_EQ(ViewExtendedTransitiveUseTime(view, /*view_color=*/5, schedule), - schedule.at(consumer)); - // Readers missing from the schedule are ignored; a view whose readers are - // all unscheduled keeps its own time. - absl::flat_hash_map only_viewbc = { - {viewbc, schedule.at(viewbc)}}; - EXPECT_EQ( - ViewExtendedTransitiveUseTime(viewbc, /*view_color=*/5, only_viewbc), - schedule.at(viewbc)); -} - -TEST_F(MemorySpaceAssignmentTest, ViewUsePrefetchDeadlineClampedToViewTime) { - // A view use extends the allocation end time through the view's transitive - // readers (`consumer` below), but a prefetched copy is materialized right - // before the view instruction itself. The prefetch deadline must therefore - // be the view's own time, not the extended reader time: with the extended - // deadline the picker may place the copy interval entirely after the - // c0/c1/c2 chain below has freed the heap, while the copy instructions - // actually run before `view`, inside the chain's live range, and the - // verifier fails with a chunk overlap. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[8]{0} parameter(0) - p1 = f32[32]{0} parameter(1) - c0 = f32[32]{0} negate(p1) - view = f32[8]{0:S(5)} custom-call(p0), custom_call_target="tpu_get_view" - viewbc = f32[8]{0:S(5)} bitcast(view) - c1 = f32[32]{0} negate(c0) - c2 = f32[32]{0} negate(c1) - cs = f32[8]{0} slice(c2), slice={[0:8]} - d0 = f32[8]{0} negate(cs) - d1 = f32[8]{0} negate(d0) - d2 = f32[8]{0} negate(d1) - consumer = f32[8]{0} add(viewbc, d2) - ROOT t = (f32[8]{0}, f32[8]{0}) tuple(consumer, d2) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.dus_view_color = 5; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker( - /*min_overlap_count=*/2, /*max_overlap_count=*/10); - ASSERT_OK(AssignMemorySpaceAndReturnStatus(module.get(), std::move(options), - /*buffer_interval_compare=*/{}, - &prefetch_interval_picker) - .status()); - // The c0/c1/c2 chain keeps the whole heap busy across `view`, so no legal - // prefetch window exists: the base must stay in default memory. - HloInstruction* view = FindInstruction(module.get(), "view"); - ASSERT_NE(view, nullptr); - EXPECT_THAT(view->operand(0), op::Parameter(0)); -} - -TEST_F(MemorySpaceAssignmentTest, - SyncDynamicSliceReplacementWithLateIndexOperand) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0_data = s32[10,2,3]{2,1,0} parameter(0) // Data for dynamic-slice - p1_long_pole = s32[10,2,3]{2,1,0} parameter(1) // For creating a long schedule gap - p2_index = s32[] parameter(2) // Index for dynamic-slice - - // Long chain of operations from p1_long_pole - negate0_lp = s32[10,2,3]{2,1,0} negate(p1_long_pole) - negate1_lp = s32[10,2,3]{2,1,0} negate(negate0_lp) - negate2_lp = s32[10,2,3]{2,1,0} negate(negate1_lp) - negate3_lp = s32[10,2,3]{2,1,0} negate(negate2_lp) - negate4_lp = s32[10,2,3]{2,1,0} negate(negate3_lp) - negate5_lp = s32[10,2,3]{2,1,0} negate(negate4_lp) - negate6_lp = s32[10,2,3]{2,1,0} negate(negate5_lp) - negate7_lp = s32[10,2,3]{2,1,0} negate(negate6_lp) // Available late in the schedule - - // Index operand, made available after the long pole computation - index = s32[] bitcast(negate7_lp) - - zero = s32[] constant(0) - - // Dynamic slice using p0_data and the late-available index - dynamic_slice = s32[1,2,3] dynamic-slice(p0_data, index, zero, zero), dynamic_slice_sizes={1,2,3} - ROOT root = s32[1,2,3] negate(dynamic_slice) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - auto module_copy = module->Clone(); - - const auto kMakeOptions = [this]() { - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; // Ensure p0_data can be prefetched - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - return options; - }; - - // The index operand is not available until the dynamic-slice instruction - // itself which pushes earliest_prefetch_time to right before the - // dynamic-slice. We test two cases with min_prefetch_interval = 2 and 1: Case - // 1 - min_prefetch_interval = 2: - // The dynamic slice is NOT replaced by an async. This is because we - // require at least 2 instructions to be overlapped (including the - // dynamic-slice). - AssignMemorySpace(module.get(), kMakeOptions(), - /*max_prefetch_interval=*/10, - /*min_prefetch_interval=*/2); - HloInstruction* root = module->entry_computation()->root_instruction(); - ASSERT_NE(root, nullptr); - ASSERT_EQ(root->operand(0)->opcode(), HloOpcode::kDynamicSlice); - - // Case 2 - min_prefetch_interval = 1: - // The dynamic slice is replaced by an async. - AssignMemorySpace(module_copy.get(), kMakeOptions(), - /*max_prefetch_interval=*/10, - /*min_prefetch_interval=*/1); - root = module_copy->entry_computation()->root_instruction(); - ASSERT_NE(root, nullptr); - ASSERT_EQ(root->operand(0)->opcode(), HloOpcode::kAsyncDone); -} - -TEST_F(MemorySpaceAssignmentTest, SyncDynamicSliceReplacementAfterPrefetch) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - index = s32[] parameter(2) - zero = s32[] constant(0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - dynamic_slice = f32[1,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={1,2,3} - concat = f32[11,2,3] concatenate(negate7, dynamic_slice), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* index = FindInstruction(module.get(), "index"); - HloInstruction* zero = FindInstruction(module.get(), "zero"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT(concat->operand(1), - op::AsyncDone(op::AsyncStart(p0, index, zero, zero))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementIgnoredTrivials) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - slice = f32[1,2,3] slice(p0), slice={[0:1], [0:2], [0:3]} - bitcast0 = f32[1,3,2] bitcast(slice) - bitcast1 = f32[10,3,2] bitcast(negate7) - concat = f32[11,3,2] concatenate(bitcast1, bitcast0), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT(concat->operand(1), - op::Bitcast(op::AsyncDone(op::AsyncStart(p0)))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncDynamicSliceReplacementIgnoredTrivials) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - index = s32[] parameter(2) - zero = s32[] constant(0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - dynamic_slice = f32[1,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={1,2,3} - bitcast0 = f32[1,3,2] bitcast(dynamic_slice) - bitcast1 = f32[10,3,2] bitcast(negate7) - concat = f32[11,3,2] concatenate(bitcast1, bitcast0), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* index = FindInstruction(module.get(), "index"); - HloInstruction* zero = FindInstruction(module.get(), "zero"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT( - concat->operand(1), - op::Bitcast(op::AsyncDone(op::AsyncStart(p0, index, zero, zero)))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementAfterEviction) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[8,4,2]{2,1,0} parameter(0) - p1 = f32[4,4,2]{2,1,0} parameter(1) - negate_p0 = negate(p0) - slice0 = f32[1,4,2] slice(negate_p0), slice={[0:1], [0:4], [0:2]} - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - negate8 = negate(negate7) - negate9 = negate(negate8) - negate10 = negate(negate9) - slice1 = f32[1,4,2] slice(negate10), slice={[0:1], [0:4], [0:2]} - concat = f32[2,4,2] concatenate(slice0, slice1), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 400; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - - AssignMemorySpace(module.get(), std::move(options)); - - HloInstruction* negate_p0 = FindInstruction(module.get(), "negate_p0"); - ASSERT_NE(negate_p0, nullptr); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT(concat->operand(0), - op::AsyncDone(op::AsyncStart(op::AsyncCopy( - kDefaultMemorySpace, kAlternateMemorySpace, negate_p0)))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncDynamicSliceReplacementAfterEviction) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[8,4,2]{2,1,0} parameter(0) - p1 = f32[4,4,2]{2,1,0} parameter(1) - index = s32[] parameter(2) - zero = s32[] constant(0) - negate_p0 = negate(p0) - dynamic_slice = f32[1,4,2] dynamic-slice(negate_p0, index, zero, zero), dynamic_slice_sizes={1,4,2} - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - negate8 = negate(negate7) - negate9 = negate(negate8) - negate10 = negate(negate9) - dynamic_slice1 = f32[1,4,2] dynamic-slice(negate10, index, zero, zero), dynamic_slice_sizes={1,4,2} - concat = f32[2,4,2] concatenate(dynamic_slice, dynamic_slice1), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 400; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* negate_p0 = FindInstruction(module.get(), "negate_p0"); - ASSERT_NE(negate_p0, nullptr); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* index = FindInstruction(module.get(), "index"); - HloInstruction* zero = FindInstruction(module.get(), "zero"); - EXPECT_THAT(concat->operand(0), - op::AsyncDone(op::AsyncStart(negate_p0, index, zero, zero))); - EXPECT_THAT(negate_p0, op::Negate(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, p0))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementTwoSlices) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - slice.1 = f32[1,2,3] slice(p0), slice={[0:1], [0:2], [0:3]} - slice.2 = f32[1,2,3] slice(p0), slice={[1:2], [0:2], [0:3]} - add = f32[1,2,3] add(slice.1, slice.2) - ROOT concat = f32[11,2,3] concatenate(negate7, add), dimensions={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* add = FindInstruction(module.get(), "add"); - ASSERT_NE(add, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - ASSERT_NE(add->operand(0), add->operand(1)); - EXPECT_THAT(add->operand(0), op::AsyncDone(op::AsyncStart(p0))); - EXPECT_THAT(add->operand(1), op::AsyncDone(op::AsyncStart(p0))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncDynamicSliceReplacementTwoSlices) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - index = s32[] parameter(2) - zero = s32[] constant(0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - dynamic_slice.1 = f32[1,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={1,2,3} - dynamic_slice.2 = f32[1,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={1,2,3} - add = f32[1,2,3] add(dynamic_slice.1, dynamic_slice.2) - ROOT concat = f32[11,2,3] concatenate(negate7, add), dimensions={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* add = FindInstruction(module.get(), "add"); - ASSERT_NE(add, nullptr); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* index = FindInstruction(module.get(), "index"); - HloInstruction* zero = FindInstruction(module.get(), "zero"); - ASSERT_NE(p0, nullptr); - ASSERT_NE(add->operand(0), add->operand(1)); - EXPECT_THAT(add->operand(0), - op::AsyncDone(op::AsyncStart(p0, index, zero, zero))); - EXPECT_THAT(add->operand(1), - op::AsyncDone(op::AsyncStart(p0, index, zero, zero))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementNestedSlices) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - slice0 = f32[9,2,3] slice(p0), slice={[0:9], [0:2], [0:3]} - negate8 = f32[10,2,3] negate(negate7) - negate9 = f32[10,2,3] negate(negate8) - negate10 = f32[10,2,3] negate(negate9) - negate11 = f32[10,2,3] negate(negate10) - negate12 = f32[10,2,3] negate(negate11) - slice1 = f32[1,2,3] slice(slice0), slice={[0:1], [0:2], [0:3]} - concat = f32[11,2,3] concatenate(negate12, slice1), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 300; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT(concat->operand(1), op::Slice(op::AsyncDone(op::AsyncStart(p0)))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncDynamicSliceReplacementNestedSlices) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - index = s32[] parameter(2) - zero = s32[] constant(0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - dynamic_slice.0 = f32[9,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={9,2,3} - negate8 = f32[10,2,3] negate(negate7) - negate9 = f32[10,2,3] negate(negate8) - negate10 = f32[10,2,3] negate(negate9) - negate11 = f32[10,2,3] negate(negate10) - negate12 = f32[10,2,3] negate(negate11) - slice1 = f32[1,2,3] slice(dynamic_slice.0), slice={[0:1], [0:2], [0:3]} - concat = f32[11,2,3] concatenate(negate12, slice1), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 300; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* index = FindInstruction(module.get(), "index"); - HloInstruction* zero = FindInstruction(module.get(), "zero"); - ASSERT_NE(p0, nullptr); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT(concat->operand(1), - op::Slice(op::AsyncDone(op::AsyncStart(p0, index, zero, zero)))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementOneFails) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[1,2,3]{2,1,0} parameter(1) - negate0 = f32[1,2,3]{2,1,0:S(1)} negate(p1) - negate1 = f32[1,2,3]{2,1,0:S(1)} negate(negate0) - negate2 = f32[1,2,3]{2,1,0:S(1)} negate(negate1) - negate3 = f32[1,2,3]{2,1,0:S(1)} negate(negate2) - negate4 = f32[1,2,3]{2,1,0:S(1)} negate(negate3) - negate5 = f32[1,2,3]{2,1,0:S(1)} negate(negate4) - negate6 = f32[1,2,3]{2,1,0:S(1)} negate(negate5) - negate7 = f32[1,2,3]{2,1,0:S(1)} negate(negate6) - slice.0 = f32[8,2,3] slice(p0), slice={[0:8], [0:2], [0:3]} - slice.1 = f32[1,2,3] slice(p0), slice={[8:9], [0:2], [0:3]} - slice.2 = f32[1,2,3] slice(p0), slice={[9:10], [0:2], [0:3]} - add.0 = f32[1,2,3] add(slice.1, slice.2) - concat.0 = f32[9,2,3] concatenate(slice.0, add.0), dimensions={0} - ROOT concat.1 = f32[10,2,3] concatenate(negate7, concat.0), dimensions={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 72; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - HloInstruction* add0 = FindInstruction(module.get(), "add.0"); - ASSERT_NE(add0, nullptr); - EXPECT_THAT(add0->operand(0), op::AsyncDone(op::AsyncStart(p0))); - EXPECT_THAT(add0->operand(1), op::AsyncDone(op::AsyncStart(p0))); - HloInstruction* slice0 = FindInstruction(module.get(), "slice.0"); - ASSERT_NE(slice0, nullptr); -} - -TEST_F(MemorySpaceAssignmentTest, SyncDynamicSliceReplacementOneFails) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[1,2,3]{2,1,0} parameter(1) - index = s32[] parameter(2) - zero = s32[] constant(0) - negate0 = f32[1,2,3]{2,1,0:S(1)} negate(p1) - negate1 = f32[1,2,3]{2,1,0:S(1)} negate(negate0) - negate2 = f32[1,2,3]{2,1,0:S(1)} negate(negate1) - negate3 = f32[1,2,3]{2,1,0:S(1)} negate(negate2) - negate4 = f32[1,2,3]{2,1,0:S(1)} negate(negate3) - negate5 = f32[1,2,3]{2,1,0:S(1)} negate(negate4) - negate6 = f32[1,2,3]{2,1,0:S(1)} negate(negate5) - negate7 = f32[1,2,3]{2,1,0:S(1)} negate(negate6) - dynamic_slice.0 = f32[8,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={8,2,3} - dynamic_slice.1 = f32[1,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={1,2,3} - dynamic_slice.2 = f32[1,2,3] dynamic-slice(p0, index, zero, zero), dynamic_slice_sizes={1,2,3} - add.0 = f32[1,2,3] add(dynamic_slice.1, dynamic_slice.2) - concat.0 = f32[9,2,3] concatenate(dynamic_slice.0, add.0), dimensions={0} - ROOT concat.1 = f32[10,2,3] concatenate(negate7, concat.0), dimensions={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 72; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - HloInstruction* index = FindInstruction(module.get(), "index"); - HloInstruction* zero = FindInstruction(module.get(), "zero"); - ASSERT_NE(p0, nullptr); - HloInstruction* add0 = FindInstruction(module.get(), "add.0"); - ASSERT_NE(add0, nullptr); - EXPECT_THAT(add0->operand(0), - op::AsyncDone(op::AsyncStart(p0, index, zero, zero))); - EXPECT_THAT(add0->operand(1), - op::AsyncDone(op::AsyncStart(p0, index, zero, zero))); - HloInstruction* slice0 = FindInstruction(module.get(), "dynamic_slice.0"); - ASSERT_NE(slice0, nullptr); -} - -// The prefetch logic has to correctly distinguish the output shape of an async -// copy vs an async slice. In this test, a prefetch of p0 would not fit into the -// memory, while prefetching a slice of p0 is feasible. -TEST_F(MemorySpaceAssignmentTest, SyncSliceReplacementTheSlicedOneFits) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[1,2,3]{2,1,0} parameter(1) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - slice = f32[1,2,3] slice(p0), slice={[0:1], [0:2], [0:3]} - concat = f32[2,2,3] concatenate(negate7, slice), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - - const auto kMakeOptions = [this](bool enable_sync_slice_replacement) { - Options options = DefaultMemorySpaceOptions(); - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - options.max_size_in_bytes = 64; - options.enable_sync_slice_replacement = enable_sync_slice_replacement; - return options; - }; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module1, - ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module1.get(), - kMakeOptions(/*enable_sync_slice_replacement=*/false)); - HloInstruction* p0 = FindInstruction(module1.get(), "p0"); - ASSERT_NE(p0, nullptr); - HloInstruction* concat = FindInstruction(module1.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT(concat->operand(1), op::Slice(p0)); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module2, - ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module2.get(), - kMakeOptions(/*enable_sync_slice_replacement=*/true)); - p0 = FindInstruction(module2.get(), "p0"); - ASSERT_NE(p0, nullptr); - concat = FindInstruction(module2.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT(concat->operand(1), op::AsyncDone(op::AsyncStart(p0))); -} - -TEST_F(MemorySpaceAssignmentTest, SyncReplacementMultipleOpTypes) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - p0_copy = copy(p0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - slice = f32[1,2,3] slice(p0_copy), slice={[0:1], [0:2], [0:3]} - concat = f32[11,2,3] concatenate(negate7, slice), dimensions={0} - ROOT root = negate(concat) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - options.enable_sync_copy_replacement = true; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* p0 = FindInstruction(module.get(), "p0"); - ASSERT_NE(p0, nullptr); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT( - concat->operand(1), - op::Slice(op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0))); -} - -// This test is for the redundant aliasing bug (b/374902759) introduced between -// different operands of the same instruction while converting sync copy to -// async ones. -TEST_F(MemorySpaceAssignmentTest, SyncReplacementAliasingBug) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true, entry_computation_layout={(f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[])->f32[10,2,3]{2,1,0}} - -%while_body (p0.1: (f32[10,2,3], f32[10,2,3], f32[10,2,3], pred[])) -> (f32[10,2,3], f32[10,2,3], f32[10,2,3], pred[]) { - %p0.1 = (f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) parameter(0) - %gte0 = f32[10,2,3]{2,1,0} get-tuple-element((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) %p0.1), index=0 - %gte1 = f32[10,2,3]{2,1,0} get-tuple-element((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) %p0.1), index=1 - %gte2 = f32[10,2,3]{2,1,0} get-tuple-element((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) %p0.1), index=2 - %gte3 = pred[] get-tuple-element((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) %p0.1), index=3 - %neg0 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %gte2) - %neg1 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg0) - %neg2 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg1) - %neg3 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg2) - %neg4 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg3) - %neg5 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg4) - %neg6 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg5) - %neg7 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %neg6) - ROOT %tuple = (f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) tuple(f32[10,2,3]{2,1,0} %gte0, f32[10,2,3]{2,1,0} %gte1, f32[10,2,3]{2,1,0} %neg7, pred[] %gte3) -} - -%while_cond (p0: (f32[10,2,3], f32[10,2,3], f32[10,2,3], pred[])) -> pred[] { - %p0 = (f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) parameter(0) - ROOT %gte = pred[] get-tuple-element((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) %p0), index=3 -} - -ENTRY %entry (p0.2: f32[10,2,3], p1: f32[10,2,3], p2: pred[]) -> f32[10,2,3] { - %p0.2 = f32[10,2,3]{2,1,0} parameter(0) - %p1 = f32[10,2,3]{2,1,0} parameter(1) - %p2 = pred[] parameter(2) - p0_copy = f32[10,2,3]{2,1,0} copy(f32[10,2,3]{2,1,0} %p0.2) - %p0_copy_copy = f32[10,2,3]{2,1,0} copy(f32[10,2,3]{2,1,0} p0_copy) - %negate0 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %p1) - %negate1 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate0) - %negate2 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate1) - %negate3 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate2) - %negate4 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate3) - %negate5 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate4) - %negate6 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate5) - %negate7 = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %negate6) - %tuple.3 = (f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) tuple(f32[10,2,3]{2,1,0} %negate7, f32[10,2,3]{2,1,0} %p0_copy, f32[10,2,3]{2,1,0} %p0_copy_copy, pred[] %p2) - while.1 = (f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) while((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) %tuple.3), condition=%while_cond, body=%while_body - %gte.1 = f32[10,2,3]{2,1,0} get-tuple-element((f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, f32[10,2,3]{2,1,0}, pred[]) while.1), index=2 - ROOT %negate = f32[10,2,3]{2,1,0} negate(f32[10,2,3]{2,1,0} %gte.1) -} - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 1024; - options.enable_sync_copy_replacement = true; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - AssignMemorySpace(module.get(), std::move(options)); - HloInstruction* while_instruction = FindInstruction(module.get(), "while.1"); - ASSERT_NE(while_instruction, nullptr); - const HloInstruction* tuple = while_instruction->operand(0); - HloInstruction* p0_copy = FindInstruction(module.get(), "p0_copy"); - ASSERT_NE(p0_copy, nullptr); - EXPECT_THAT(tuple->operand(1), op::AsyncCopy(kDefaultMemorySpace, - kAlternateMemorySpace, p0_copy)); - EXPECT_THAT(tuple->operand(2), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - tuple->operand(1))); -} - -// Added for b/376344953 that was introduced when we tried to -// convert a sync copy that was used by a conditional into an async copy. -TEST_F(MemorySpaceAssignmentTest, ConditionalCopyReplacement) { - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg1 = f32[3]{0} negate(gte) - } - - false_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg2 = f32[3]{0} negate(gte) - } - - ENTRY entry { - p0_main = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0_main) - tuple = (f32[3]{0}) tuple(copy) - ROOT conditional = f32[3]{0} conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_sync_copy_replacement = true; - AssignMemorySpace(module.get(), std::move(options)); - auto conditional = - module->GetComputationWithName("entry")->GetInstructionWithName( - "conditional"); - CHECK_NE(conditional, nullptr); - auto p0 = module->GetComputationWithName("entry")->GetInstructionWithName( - "p0_main"); - CHECK_NE(p0, nullptr); - auto copy = conditional->operand(1)->operand(0); - CHECK_NE(copy, nullptr); - EXPECT_THAT(copy, - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); -} - -TEST_F(MemorySpaceAssignmentTest, AllocationRequestAndResultModifierTest) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = f32[2,3]{1,0} negate(p1) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(negate2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[2,3]{1,0} negate(negate5) - negate7 = f32[2,3]{1,0} negate(negate6) - ROOT add0 = f32[2,3]{1,0} add(p0, negate7) - } - )hlo"; - // The baseline behavior is to prefetch p0 at add0. - ASSERT_OK_AND_ASSIGN(std::unique_ptr baseline_module, - ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(baseline_module.get(), DefaultMemorySpaceOptions()); - HloInstruction* add0 = FindInstruction(baseline_module.get(), "add0"); - ASSERT_NE(add0, nullptr); - HloInstruction* p0 = FindInstruction(baseline_module.get(), "p0"); - ASSERT_NE(p0, nullptr); - EXPECT_THAT(add0->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - - // We should be able to prevent prefetching p0 at add0 using - // allocation_result_modifier_testing_fn. - ASSERT_OK_AND_ASSIGN( - std::unique_ptr result_modifier_module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options_result_modifier = DefaultMemorySpaceOptions(); - options_result_modifier.max_retries = 1; - options_result_modifier.allocation_request_modifier_testing_fn = nullptr; - options_result_modifier.allocation_result_modifier_testing_fn = - [](const AllocationRequest& request, AllocationResult& result, - int64_t retry_number) { - if (request.allocation_value_to_update->position() - .instruction->name() == "p0" && - request.use->hlo_use.instruction->name() == "add0") { - result = AllocationResult::kFailRequiresUncommit; - } - }; - AssignMemorySpace(result_modifier_module.get(), - std::move(options_result_modifier)); - add0 = FindInstruction(result_modifier_module.get(), "add0"); - ASSERT_NE(add0, nullptr); - p0 = FindInstruction(result_modifier_module.get(), "p0"); - ASSERT_NE(p0, nullptr); - EXPECT_EQ(add0->operand(0), p0); - - // We should be able to enforce an earlier prefetch of p0 at add0 using - // allocation_request_modifier_testing_fn. - ASSERT_OK_AND_ASSIGN( - std::unique_ptr request_modifier_module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options_request_modifier = DefaultMemorySpaceOptions(); - options_request_modifier.max_retries = 1; - options_request_modifier - .allocation_request_modifier_testing_fn = [](AllocationRequest& request) { - if (request.allocation_value_to_update->position().instruction->name() == - "p0" && - request.use->hlo_use.instruction->name() == "add0") { - // Schedule the copy-done before negate4 (scheduled at 6). - request.latest_prefetch_time = 6; - } - }; - options_request_modifier.allocation_result_modifier_testing_fn = nullptr; - AssignMemorySpace(request_modifier_module.get(), - std::move(options_request_modifier)); - add0 = FindInstruction(request_modifier_module.get(), "add0"); - CHECK_NE(add0, nullptr); - p0 = FindInstruction(request_modifier_module.get(), "p0"); - CHECK_NE(p0, nullptr); - EXPECT_THAT(add0->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, p0)); - // The copy-done should have been scheduled before negate4. - HloInstruction* negate4 = - FindInstruction(request_modifier_module.get(), "negate4"); - CHECK_NE(negate4, nullptr); - const HloInstructionSequence& sequence = - request_modifier_module->schedule().sequence( - request_modifier_module->entry_computation()); - auto find_index = [&](const HloInstruction* instruction) { - return std::distance(sequence.instructions().begin(), - absl::c_find(sequence.instructions(), instruction)); - }; - - int negate4_index = find_index(negate4); - int copy_done_index = find_index(add0->operand(0)); - EXPECT_LT(copy_done_index, negate4_index); -} - -// Added for b/372277844#comment15 that was introduced when the allocation -// failed while trying to convert a sync slice to an async one, but not due to -// the conversion itself. In this case, associated buffer with the slice -// (p0_copy) is too large to fit in alternate memory. Hence, the -// allocation_values will be empty in retries, previously causing a crash in -// MsaAlgorithm::GetInefficientAllocationSites(). -TEST_F(MemorySpaceAssignmentTest, SyncReplacementLargeBuffers) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[10,2,3]{2,1,0} parameter(0) - p1 = f32[10,2,3]{2,1,0} parameter(1) - p0_copy = f32[10,2,3]{2,1,0} copy(p0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - negate3 = negate(negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - slice = f32[1,2,3] slice(p0_copy), slice={[0:1], [0:2], [0:3]} - ROOT concat = f32[11,2,3] concatenate(negate5, slice), dimensions={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 64; - options.max_retries = 2; - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - // Force the allocation of p0_copy to fail for the concat use with only - // AllocationResult::kFailRequiresUncommit. This means that while the slice - // replacement was successful, the algorithm must retry one more time without - // sync slice conversion target, so that maybe other constraints of the - // allocation can be satisfied. - options.allocation_result_modifier_testing_fn = - [](const AllocationRequest& request, AllocationResult& result, - int64_t retry_number) { - if (request.allocation_value->position().instruction->name() == - "p0_copy" && - request.use->hlo_use.instruction->name() == "concat") { - result = AllocationResult::kFailRequiresUncommit; - } - }; - // options.inefficient_use_to_copy_ratio must be greater than 0 and the cost - // model must be set to trigger the inefficient allocation site logic. - options.inefficient_use_to_copy_ratio = 1.0; - AssignMemorySpaceUsingCostAnalysis(module.get(), std::move(options)); - - HloInstruction* p0_copy = FindInstruction(module.get(), "p0_copy"); - ASSERT_NE(p0_copy, nullptr); - HloInstruction* concat = FindInstruction(module.get(), "concat"); - ASSERT_NE(concat, nullptr); - EXPECT_THAT(concat->operand(1), op::Slice(p0_copy)); -} - -// Added for b/376869021, which surfaced when we tried to convert a sync slice -// that had to extend the allocation of its operand in the alternate memory. In -// this test, we expect the slice0 operand (p0_copy) maintain a valid allocation -// in the alternate memory, until it gets transferred by the async replacement -// of slice0. We hence stress-test such validity by delaying the allocation of -// slice0 by 3 steps. -TEST_F(MemorySpaceAssignmentTest, SyncReplacementAllocationExtensionBug) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,2,3]{2,1,0} parameter(0) - p1 = f32[4,2,3]{2,1,0} parameter(1) - p0_copy = f32[2,2,3]{2,1,0} copy(p0) - negate0 = negate(p1) - negate1 = negate(negate0) - negate2 = negate(negate1) - p0_copy0_negate = negate(p0_copy) - copy_negate2 = copy(negate2) - slice0 = f32[1,2,3] slice(p0_copy), slice={[0:1], [0:2], [0:3]} - negate3 = negate(copy_negate2) - negate4 = negate(negate3) - negate5 = negate(negate4) - negate6 = negate(negate5) - negate7 = negate(negate6) - neg_slice0 = negate(slice0) - ROOT tuple = tuple(negate7, neg_slice0) -})hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_sync_copy_replacement = false; - options.enable_sync_slice_replacement = true; - options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { return true; }; - options.max_size_in_bytes = 96; - options.is_position_allowed_in_alternate_mem_fn = - [](const HloPosition& position) { - return position.instruction->name() != "p0_copy"; - }; - // Delay the allocation of slice0 by 3 steps to allow copy_negate2 to be - // allocated in alternate memory. - options.allocation_request_modifier_testing_fn = - [](AllocationRequest& request) { - if (request.only_extend_existing_allocation) { - request.inclusive_start_time += 3; - request.end_time += 3; - } - }; - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "copy_negate2|p0_copy" } - override_options { assign_first: true } - })pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - auto preset_assignments = AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(options), - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - HloInstruction* p0_copy = FindInstruction(module.get(), "p0_copy"); - ASSERT_NE(p0_copy, nullptr); - HloInstruction* neg_slice0 = FindInstruction(module.get(), "neg_slice0"); - ASSERT_NE(neg_slice0, nullptr); - EXPECT_THAT(neg_slice0->operand(0), op::AsyncDone(op::AsyncStart(p0_copy))); -} - -TEST_F(MemorySpaceAssignmentTest, AlwaysSpillJitPrefetchTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - // For buffers produced in alternate memory spill to default and prefetch - // just in time for uses other than immediate use (if any) and make all - // prefetches single use for first use and create new prefetches for all - // subsequent uses. - - // We expect MSA to start prefetching p1 immediately the parameter(1) - // instruction and to finish immediately before add. The - // always_spill_to_default_memory option will move the start of the prefetch - // from just after parameter(1) to just before its completion. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(negate2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[2,3]{1,0} negate(negate5) - ROOT add = f32[2,3]{1,0} add(negate6, p1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.always_spill_to_default_memory = true; - AssignMemorySpace(module.get(), std::move(options)); - const HloInstructionSequence& sequence = - module->schedule().sequence(module->entry_computation()); - for (int i = 0; i < sequence.instructions().size(); ++i) { - VLOG(2) << i << " " << sequence.instructions()[i]->ToString(); - } - ASSERT_OK_AND_ASSIGN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - const HloInstruction* add = FindInstruction(module.get(), "add"); - const HloInstruction* cd = add->operand(1); - // Check copy made just in time for use and copy is a prefetch. - EXPECT_THAT(cd, op::CopyDone()); - EXPECT_EQ(live_range->instruction_schedule().at(add), - live_range->instruction_schedule().at(cd) + 1); - const HloInstruction* cs = cd->operand(0); - EXPECT_THAT(cs, op::CopyStart()); - EXPECT_EQ(live_range->instruction_schedule().at(add), - live_range->instruction_schedule().at(cs) + 2); - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); -} - -TEST_F(MemorySpaceAssignmentTest, AlwaysSpillPrefetchForSecondUseTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - // - // Setting always_spill_to_default_memory option to true makes sure the - // negate0 buffer is copied to default memory between negate0 and negate1, - // so that version can be prefetched just before it is used at add0. - // Additionally, we leave a copy of negate0 in alternate memory for use at - // negate1. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(negate2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - add0 = f32[2,3]{1,0} add(negate5, negate0) - ROOT add1 = f32[2,3]{1,0} add(add0, p1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.always_spill_to_default_memory = true; - AssignMemorySpace(module.get(), std::move(options)); - const HloInstructionSequence& sequence = - module->schedule().sequence(module->entry_computation()); - for (int i = 0; i < sequence.instructions().size(); ++i) { - VLOG(2) << i << " " << sequence.instructions()[i]->ToString(); - } - ASSERT_OK_AND_ASSIGN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - // Check copies are made just in time for use and copies are prefetches. - const HloInstruction* add1 = FindInstruction(module.get(), "add1"); - const HloInstruction* cd1 = add1->operand(1); - EXPECT_THAT(cd1, op::CopyDone()); - EXPECT_EQ(live_range->instruction_schedule().at(add1), - live_range->instruction_schedule().at(cd1) + 1); - const HloInstruction* cs1 = cd1->operand(0); - EXPECT_THAT(cs1, op::CopyStart()); - EXPECT_EQ(live_range->instruction_schedule().at(add1), - live_range->instruction_schedule().at(cs1) + 2); - EXPECT_EQ(cd1->shape().layout().memory_space(), kAlternateMemorySpace); - const HloInstruction* add0 = FindInstruction(module.get(), "add0"); - - const HloInstruction* cd0 = add0->operand(1); - EXPECT_THAT(cd0, op::CopyDone()); - EXPECT_EQ(live_range->instruction_schedule().at(add0), - live_range->instruction_schedule().at(cd0) + 1); - const HloInstruction* cs0 = cd0->operand(0); - EXPECT_THAT(cs0, op::CopyStart()); - EXPECT_EQ(live_range->instruction_schedule().at(add0), - live_range->instruction_schedule().at(cs0) + 2); - EXPECT_EQ(cd0->shape().layout().memory_space(), kAlternateMemorySpace); - // Check prefetch was made from an eviction. - const HloInstruction* eviction_done = cs0->operand(0); - EXPECT_EQ(eviction_done->shape().layout().memory_space(), - kDefaultMemorySpace); - const HloInstruction* evection_start = eviction_done->operand(0); - const HloInstruction* negate0 = evection_start->operand(0); - // Check eviction was immediate. - EXPECT_EQ(live_range->instruction_schedule().at(evection_start), - live_range->instruction_schedule().at(negate0) + 1); - EXPECT_EQ(live_range->instruction_schedule().at(eviction_done), - live_range->instruction_schedule().at(negate0) + 2); - EXPECT_EQ(negate0->name(), "negate0"); -} - -TEST_F(MemorySpaceAssignmentTest, AlwaysSpillEvictionTest) { - // tanh0 buffer is produced in alternate memory and it has two uses that are - // sufficiently far apart for an eviction to be scheduled. When the - // always_spill_to_default_memory option is not true, the buffer stays in - // alternate memory to serve the first use, is evicted and prefetched again - // for second use. Setting always_spill_to_default_memory option to true makes - // the eviction immediate, right after tanh0, the first use at add5 and second - // use at tuple are served from separate, just-in-time prefetches that copy - // from the eviction that previously occurred. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - tanh0 = f32[4,3]{1,0} tanh(p0) - add0 = f32[4,3]{1,0} add(p0, p0) - add1 = f32[4,3]{1,0} add(add0, p0) - add2 = f32[4,3]{1,0} add(add1, p0) - add3 = f32[4,3]{1,0} add(add2, p0) - add4 = f32[4,3]{1,0} add(add3, p0) - add5 = f32[4,3]{1,0} add(add4, tanh0) - negate0 = f32[4,3]{1,0} negate(add5) - tanh1 = f32[4,3]{1,0} tanh(negate0) - negate1 = f32[4,3]{1,0} negate(negate0) - tanh2 = f32[4,3]{1,0} tanh(tanh1) - negate2 = f32[4,3]{1,0} negate(negate1) - ROOT tuple = tuple(tanh0, tanh2, negate2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.always_spill_to_default_memory = true; - AssignMemorySpace(module.get(), std::move(options)); - const HloInstructionSequence& sequence = - module->schedule().sequence(module->entry_computation()); - for (int i = 0; i < sequence.instructions().size(); ++i) { - VLOG(2) << i << " " << sequence.instructions()[i]->ToString(); - } - ASSERT_OK_AND_ASSIGN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(std::unique_ptr live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - // 1. Check tanh0 buffer is short lived. - // 2. Check tanh0 eviction is immediate. - // 3. Check tuple is served from eviction. - // 4. Check add5 is served from a prefetch. - // 5. Check prefetch comes from the immediate eviction. - const HloInstruction* tuple = FindInstruction(module.get(), "tuple"); - const HloInstruction* tanh0_eviction_done = tuple->operand(0); - const HloInstruction* tanh0_eviction_start = tanh0_eviction_done->operand(0); - const HloInstruction* tanh0 = tanh0_eviction_start->operand(0); - EXPECT_EQ(tanh0->name(), "tanh0"); - EXPECT_EQ(tanh0_eviction_done->shape().layout().memory_space(), - kDefaultMemorySpace); - EXPECT_EQ(live_range->instruction_schedule().at(tanh0_eviction_start), - live_range->instruction_schedule().at(tanh0) + 1); - EXPECT_EQ(live_range->instruction_schedule().at(tanh0_eviction_done), - live_range->instruction_schedule().at(tanh0) + 2); - const HloInstruction* add5 = FindInstruction(module.get(), "add5"); - const HloInstruction* tanh0_prefetch_done = add5->operand(1); - const HloInstruction* tanh0_prefetch_start = tanh0_prefetch_done->operand(0); - EXPECT_EQ(tanh0_prefetch_done->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(live_range->instruction_schedule().at(add5), - live_range->instruction_schedule().at(tanh0_prefetch_done) + 1); - EXPECT_EQ(live_range->instruction_schedule().at(add5), - live_range->instruction_schedule().at(tanh0_prefetch_start) + 2); - EXPECT_EQ(tanh0_eviction_done, tanh0_prefetch_start->operand(0)); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdatePreferredPrefetchTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { size_lte: 24 size_gte: 24 } - override_options { prefetch_eagerness: 0.5 } - })pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdateConfigExactMatchBeforeTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: "add" operand_number: 1 } - override_options { - before_instruction: { instruction_regex: "%?negate.3 =.*" } - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdateConfigExactMatchAfterTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: "add" operand_number: 1 } - override_options { - after_instruction: { instruction_regex: "%?negate.1 =.*" } - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[4], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdateConfigExactMatchTooLateTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: "add" operand_number: 1 } - override_options { - after_instruction: { instruction_name_regex: "%?negate.5" } - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - // Ensure the Async copy is not scheduled. - EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdateConfigPrecedenceTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { size_lte: 24 size_gte: 24 } - override_options { prefetch_eagerness: 0.5 } - } - overrides { - hlo_operand_filter { instruction_name_regex: "add" operand_number: 1 } - override_options { - after_instruction: { instruction_name_regex: "%?negate.1" } - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdateConfigExactMatchPrecedenceTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: "add" operand_number: 1 } - override_options { - after_instruction: { instruction_name_regex: "%?negate.1" } - } - } - overrides { - hlo_operand_filter { size_lte: 24 size_gte: 24 } - override_options { prefetch_eagerness: 0.5 } - } - )pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[4], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, FilterUpdatePreferredPrefetchNoMatchTest) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { size_lte: 24 size_gte: 25 } - override_options { prefetch_eagerness: 0.5 } - } - )pb"; - ASSERT_OK_AND_ASSIGN(options.preferred_prefetch_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); - // Ensure the CopyStart/CopyDone schedules. - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[2], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPrefetchExactTimeSuccess) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - prefetch { - after_instruction { instruction_name_regex: ".*negate.3.*" } - strict_timing: true - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); - // Ensure the CopyStart is scheduled after negate.3 (schedule index 6). - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); - EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); - EXPECT_THAT(sequence.instructions()[5], op::Negate()); - EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); - EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPrefetchLogicalTimeSuccess) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - prefetch { logical_time: 4 strict_timing: true } - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPrefetchUnsatisfiableFails) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate1, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - prefetch { logical_time: 0 strict_timing: true } - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - InstructionHoister instruction_hoister; - CHECK_OK(instruction_hoister.Run(module.get()).status()); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - auto status_or = AssignMemorySpaceAndReturnStatus( - module.get(), std::move(options), /*buffer_interval_compare=*/{}, - &prefetch_interval_picker); - EXPECT_FALSE(status_or.ok()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPinInAlternateMemory) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate1, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: ".*p1.*" } - pin_in_alternate_memory {} - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - // p1 is pinned in alternate memory, so no copy allocation is created and - // add consumes p1 directly. - EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); - EXPECT_EQ(p1->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, StrictKeepInDefaultMemory) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - // Strictly keep p1 in default memory, preventing prefetch. - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - keep_in_default_memory {} - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - // Without override, p1 would be prefetched via AsyncCopy. - // With keep_in_default_memory, add consumes Parameter(1) directly. - EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); - EXPECT_EQ(p1->shape().layout().memory_space(), kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, StrictKeepInDefaultMemoryDefiningPosition) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: ".*p1.*" } - keep_in_default_memory {} - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); - EXPECT_EQ(p1->shape().layout().memory_space(), kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - PositionMatcherDoesNotMatchConsumerInstruction) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - // Matching ".*add.*" with HloPositionMatcher should ONLY match buffers - // PRODUCED by add, NOT buffers CONSUMED by add (like p1). - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: ".*add.*" } - keep_in_default_memory {} - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - // Since p1's defining position is p1 (not add), it is NOT kept in default - // memory and gets prefetched into alternate memory space. - EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(1)))); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPrefetchBeforeInstructionSuccess) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - prefetch { - before_instruction { instruction_name_regex: ".*negate.4.*" } - strict_timing: true - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPrefetchExactTimeWithStrictBool) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - prefetch { logical_time: 4 strict_timing: true } - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPrefetchCostAnalysisExactTimeSuccess) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - - const std::string text_proto = R"pb( - overrides { - hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } - prefetch { - after_instruction { instruction_name_regex: ".*negate.3.*" } - strict_timing: true - } - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpaceUsingCostAnalysis(module.get(), std::move(options)); - - const HloInstructionSequence& sequence = - module->schedule().sequence(computation); - EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); -} - -TEST_F(MemorySpaceAssignmentTest, StrictPinInAlternateMemoryFailsOnOom) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); // 24 bytes - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, negate0}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 10; - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: ".*p0.*" } - pin_in_alternate_memory {} - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, - ParseTextProto(text_proto)); - - InstructionHoister instruction_hoister; - CHECK_OK(instruction_hoister.Run(module.get()).status()); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - auto status_or = AssignMemorySpaceAndReturnStatus( - module.get(), std::move(options), /*buffer_interval_compare=*/{}, - &prefetch_interval_picker); - EXPECT_FALSE(status_or.ok()); -} - -TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetch) { - std::unique_ptr module = CreateEvictAndPrefetchModule(); - - AssignMemorySpace(module.get()); - - EXPECT_THAT( - module->entry_computation()->root_instruction(), - op::Add(op::Add(), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::AsyncCopy(kDefaultMemorySpace, - kAlternateMemorySpace, op::Tanh())))); -} - -TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetchLimitAsyncCopies0) { - std::unique_ptr module = CreateEvictAndPrefetchModule(); - - AssignMemorySpace(module.get(), - UpdateMaxAsyncCopies(DefaultMemorySpaceOptions(), 0)); - - EXPECT_LE(CountMaximumOutstandingAsyncCopies(*module).max_prefetches, 0); - EXPECT_LE(CountMaximumOutstandingAsyncCopies(*module).max_evictions, 0); -} - -TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetchLimitAsyncCopies1) { - std::unique_ptr module = CreateEvictAndPrefetchModule(); - - AssignMemorySpace(module.get(), - UpdateMaxAsyncCopies(DefaultMemorySpaceOptions(), 1)); - - EXPECT_LE(CountMaximumOutstandingAsyncCopies(*module).max_prefetches, 1); - EXPECT_LE(CountMaximumOutstandingAsyncCopies(*module).max_evictions, 1); -} - -TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetchLimitAsyncCopies2) { - std::unique_ptr module = CreateEvictAndPrefetchModule(); - - AssignMemorySpace(module.get(), - UpdateMaxAsyncCopies(DefaultMemorySpaceOptions(), 2)); - - EXPECT_LE(CountMaximumOutstandingAsyncCopies(*module).max_prefetches, 2); - EXPECT_LE(CountMaximumOutstandingAsyncCopies(*module).max_evictions, 2); -} - -// TODO(berkin): This test is broken with some prefetch timing improvements. -TEST_F(MemorySpaceAssignmentTest, - DISABLED_DontEvictWhenThereIsDefaultMemAllocation) { - // This test is the same as EvictAndPrefetchLimitAsyncCopies1, except we check - // that there is no eviction if not necessary (due to an existing allocation - // in default memory). - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* tanh = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - // tanh should be placed in the alternate memory since there isn't much - // contention in the beginning. However, tanh has another consumer at the end. - // So it should be kicked out to default memory and prefetched back in. The - // graph below is meant to increase the contention to force eviction/prefetch - // behavior. - HloInstruction* a = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, tanh)); - HloInstruction* b = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* c = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, p0, p1)); - HloInstruction* d = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* e = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, b)); - HloInstruction* f = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, c)); - HloInstruction* g = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, d)); - HloInstruction* h = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, b, c)); - HloInstruction* i = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, b, d)); - HloInstruction* j = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, c, d)); - HloInstruction* k = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, e, f)); - HloInstruction* l = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, g, h)); - HloInstruction* m = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, i, j)); - HloInstruction* n = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, k, l)); - HloInstruction* o = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, n, m)); - // tanh is being used at the root instruction, and this should be - // prefetched. - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, o, tanh)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, tanh, a, b, c, d, e, f, g, h, i, - j, k, l, m, n, o, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), - UpdateMaxAsyncCopies(DefaultMemorySpaceOptions(), 1)); - - // We expect the second argument to multiply is prefetched c. - EXPECT_THAT(f, op::Multiply(op::Add(), op::CopyDone())); - // We make sure that the second argument to this multiply is not evicted - // CopyDone but is the original c. - EXPECT_THAT(h, op::Multiply(op::Subtract(), op::Multiply())); -} - -TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetchAndPrefetch) { - // Test for a memory corruption bug involving evict/prefetch/prefetch pattern, - // where the last prefetch copied from the original buffer in alternate buffer - // instead of evicted buffer. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* tanh = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* a = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, tanh)); - HloInstruction* b = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* c = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, p0, p1)); - HloInstruction* d = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* e = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, b)); - HloInstruction* f = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, c)); - HloInstruction* g = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, d)); - HloInstruction* h = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, b, c)); - HloInstruction* i = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, b, d)); - HloInstruction* j = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, c, d)); - HloInstruction* k = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, e, f)); - HloInstruction* l = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, g, h)); - HloInstruction* m = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, i, j)); - HloInstruction* n = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, k, l)); - HloInstruction* o = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, n, m)); - HloInstruction* add0 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, o, tanh)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, add0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* negate8 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate7)); - HloInstruction* negate9 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate8)); - HloInstruction* add1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate9, tanh)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, - {p0, p1, tanh, a, b, c, d, e, - f, g, h, i, j, k, l, m, - n, o, add0, negate0, negate1, negate2, negate3, negate4, - negate5, negate6, negate7, negate8, negate9, add1}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - // Check that both prefetches (add0 and add1) prefetch from the eviction - // instead of tanh, which will be placed in the alternate memory directly. - EXPECT_THAT( - add0, - op::Add(op::Add(), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::AsyncCopy(kDefaultMemorySpace, - kAlternateMemorySpace, op::Tanh())))); - EXPECT_THAT( - add1, - op::Add(op::Negate(), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::AsyncCopy(kDefaultMemorySpace, - kAlternateMemorySpace, op::Tanh())))); -} - -TEST_F(MemorySpaceAssignmentTest, While) { - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, scalar_shape}); - - auto cond_builder = HloComputation::Builder("WhileCond"); - // Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element) - HloInstruction* cond_param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "cond_param")); - HloInstruction* cond_iter = cond_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1)); - HloInstruction* cond_limit = cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(50.f))); - // Free cond_param[] (16 bytes), Alloc PRED[] (1 byte) - HloInstruction* cond_lt = cond_builder.AddInstruction( - HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), cond_iter, - cond_limit, ComparisonDirection::kLt)); - HloComputation* cond_computation = - module->AddEmbeddedComputation(cond_builder.Build()); - - auto body_builder = HloComputation::Builder("WhileBody"); - // Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element) - HloInstruction* body_param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "body_param")); - HloInstruction* body_iter = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, body_param, 1)); - HloInstruction* body_data = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, body_param, 0)); - HloInstruction* body_iter_increment = body_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.f))); - HloInstruction* body_iter_next = - body_builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape, HloOpcode::kAdd, body_iter, body_iter_increment)); - HloInstruction* body_data_increment = - body_builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.f, 2.f, 3.f}, {4.f, 5.f, 6.f}}))); - HloInstruction* body_data_mul = - body_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kMultiply, body_data, body_data)); - HloInstruction* body_data_add = - body_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, body_data, body_data_increment)); - HloInstruction* body_data_next = - body_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, body_data_add, body_data_mul)); - HloInstruction* body_out = body_builder.AddInstruction( - HloInstruction::CreateTuple({body_data_next, body_iter_next})); - HloComputation* body_computation = - module->AddEmbeddedComputation(body_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* data = builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "param_iter")); - HloInstruction* iter = builder.AddInstruction( - HloInstruction::CreateParameter(1, scalar_shape, "param_data")); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({data, iter})); - HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( - tuple_shape, cond_computation, body_computation, tuple)); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(cond_computation, - {cond_param, cond_iter, cond_limit, cond_lt}); - schedule.set_sequence(body_computation, - {body_param, body_iter, body_data, body_iter_increment, - body_iter_next, body_data_increment, body_data_mul, - body_data_add, body_data_next, body_out}); - schedule.set_sequence(entry_computation, {iter, data, tuple, while_op}); - CHECK_OK(module->set_schedule(schedule)); - - LOG(INFO) << module->ToString(HloPrintOptions::ShortParsable()); - - AssignMemorySpace(module.get()); - - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - EXPECT_THAT(body_data_mul, op::ShapeWithLayout(shape_in_alternate_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, Tuple) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape inner_tuple_shape = ShapeUtil::MakeTupleShape({shape}); - Shape tuple_shape = - ShapeUtil::MakeTupleShape({shape, shape, inner_tuple_shape}); - HloInstruction* p = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* p0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 0)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 1)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - HloInstruction* p2 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(inner_tuple_shape, p, 2)); - HloInstruction* p2_0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p2, 0)); - HloInstruction* mul = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, add, p2_0)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, {p, p0, negate0, negate1, negate2, negate3, negate4, negate5, - negate6, p1, add, p2, p2_0, mul}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - EXPECT_THAT( - mul, - op::Multiply(op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::GetTupleElement())), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::GetTupleElement())))); -} - -TEST_F(MemorySpaceAssignmentTest, Bitcast) { - // Bitcasts can cause the position in the alternate memory to appear multiple - // times in the preset assignments. This test ensure the preset assignments - // refer to unique positions. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateParameter(1, param_shape, "p1")); - HloInstruction* negate = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* bitcast = builder.AddInstruction( - HloInstruction::CreateBitcast(param_shape, negate)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(param_shape, HloOpcode::kAdd, bitcast, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate, bitcast, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - bitcast = add->mutable_operand(0); - EXPECT_EQ(bitcast->opcode(), HloOpcode::kBitcast); - EXPECT_EQ(bitcast->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, Bitcast2) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateParameter(1, param_shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(shape, p1)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, bitcast, negate4)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, bitcast, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - EXPECT_EQ(add->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, Bitcast3) { - HloComputation::Builder builder(TestName()); - Shape shape1 = ShapeUtil::MakeShape(F32, {2, 3}); - Shape shape2 = ShapeUtil::MakeShape(F32, {3, 2}); - Shape shape3 = ShapeUtil::MakeShape(F32, {1, 6}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape1, "p0")); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateParameter(1, param_shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape1, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape1, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape1, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape1, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape1, HloOpcode::kNegate, negate3)); - HloInstruction* bitcast1 = - builder.AddInstruction(HloInstruction::CreateBitcast(shape1, p1)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape1, HloOpcode::kAdd, bitcast1, negate4)); - HloInstruction* bitcast2 = - builder.AddInstruction(HloInstruction::CreateBitcast(shape3, p1)); - HloInstruction* bitcast3 = - builder.AddInstruction(HloInstruction::CreateBitcast(shape2, bitcast2)); - HloInstruction* bitcast4 = - builder.AddInstruction(HloInstruction::CreateBitcast(shape2, add)); - HloInstruction* mul = builder.AddInstruction(HloInstruction::CreateBinary( - shape2, HloOpcode::kMultiply, bitcast3, bitcast4)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, - {p0, p1, negate0, negate1, negate2, negate3, negate4, - bitcast1, add, bitcast2, bitcast3, bitcast4, mul}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - // We expect one bitcast on the LHS of multiply since bitcast(bitcast(foo)) is - // converted to bitcast(foo). - EXPECT_THAT( - mul, - op::Multiply( - op::Bitcast(op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::Parameter(1))), - op::Bitcast(op::Add( - op::Bitcast(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, op::Parameter(1))), - op::Negate())))); - EXPECT_EQ(add->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(add->shape().layout().memory_space(), kAlternateMemorySpace); - // bitcast2 will no longer have a consumer and should get DCE'd, so we don't - // care about its memory space. - EXPECT_EQ(mul->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(mul->operand(1)->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, BitcastTuple) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - - auto module = CreateNewVerifiedModule(); - HloComputation::Builder fusion_builder("fusion"); - HloInstruction* fusion_param = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* fusion_element0 = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion_param, 0)); - HloInstruction* fusion_element1 = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion_param, 1)); - fusion_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, fusion_element0, fusion_element1)); - HloComputation* fusion_computation = - module->AddEmbeddedComputation(fusion_builder.Build()); - - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateParameter(1, param_shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(shape, p1)); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({bitcast, p0})); - HloInstruction* fusion = builder.AddInstruction(HloInstruction::CreateFusion( - shape, HloInstruction::FusionKind::kCustom, {tuple}, fusion_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, - {p0, p1, negate0, negate1, negate2, negate3, negate4, - bitcast, tuple, fusion}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, BitcastGetTupleElementTuple) { - // This test pattern was encountered in - // //third_party/tensorflow/compiler/xla/tests:slice_test and was causing a - // breakage when there is a GetTupleElement(Tuple(Bitcast())) pattern. Also - // added a GetTupleElement(GetTupleElement(Tuple(Tuple(Bitcast())))) pattern. - absl::string_view hlo_string = R"hlo( - HloModule DoIt_S64_10_0_5_1.3, is_scheduled=true - - ENTRY %DoIt_S64_10_0_5_1.3 (p0.1: (u32[10], u32[10])) -> (u32[5], u32[5]) { - %p0.1 = (u32[10]{0:T(128)}, u32[10]{0:T(128)}) parameter(0) - %get-tuple-element.1 = u32[10]{0:T(128)} get-tuple-element((u32[10]{0:T(128)}, u32[10]{0:T(128)}) %p0.1), index=1 - %bitcast.1 = u32[5]{0:T(128)} bitcast(u32[10]{0:T(128)} %get-tuple-element.1) - %get-tuple-element = u32[10]{0:T(128)} get-tuple-element((u32[10]{0:T(128)}, u32[10]{0:T(128)}) %p0.1), index=0 - %bitcast = u32[5]{0:T(128)} bitcast(u32[10]{0:T(128)} %get-tuple-element) - %tuple.1 = (u32[5]{0:T(128)}, u32[5]{0:T(128)}) tuple(u32[5]{0:T(128)} %bitcast, u32[5]{0:T(128)} %bitcast.1) - %tuple.3 = ((u32[5]{0:T(128)}, u32[5]{0:T(128)}), (u32[5]{0:T(128)}, u32[5]{0:T(128)})) tuple(%tuple.1, %tuple.1) - %get-tuple-element.4 = u32[5]{0:T(128)} get-tuple-element((u32[5]{0:T(128)}, u32[5]{0:T(128)}) %tuple.1), index=0 - %get-tuple-element.5 = (u32[5]{0:T(128)}, u32[5]{0:T(128)}) get-tuple-element(%tuple.3), index=0 - %get-tuple-element.6 = u32[5]{0:T(128)} get-tuple-element((u32[5]{0:T(128)}, u32[5]{0:T(128)}) %get-tuple-element.5), index=1 - %copy.2 = u32[5]{0:T(128)} copy(u32[5]{0:T(128)} %get-tuple-element.4) - %copy.3 = u32[5]{0:T(128)} copy(u32[5]{0:T(128)} %get-tuple-element.6) - ROOT %tuple.2 = (u32[5]{0:T(128)}, u32[5]{0:T(128)}) tuple(u32[5]{0:T(128)} %copy.2, u32[5]{0:T(128)} %copy.3) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, GetSimplifiedOperandBug) { - // Test case for a bug finding Bitcasts in GTE(Tuple(...)) pattern. - absl::string_view hlo_string = R"hlo( - HloModule sort.16, is_scheduled=true - - ENTRY %sort.16 (param.0.1: s32[1], param.1.2: f32[1], param.2.3: u32[1], param.3.4: s32[1]) -> (s32[1], f32[1], u32[1], s32[1]) { - %param.3.4 = s32[1]{0:T(128)} parameter(3) - %param.2.3 = u32[1]{0:T(128)} parameter(2) - %param.1.2 = f32[1]{0:T(128)} parameter(1) - %param.0.1 = s32[1]{0:T(128)} parameter(0) - %tuple.1 = (s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) tuple(s32[1]{0:T(128)} %param.0.1, f32[1]{0:T(128)} %param.1.2, u32[1]{0:T(128)} %param.2.3, s32[1]{0:T(128)} %param.3.4) - %get-tuple-element.4 = s32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=0 - %get-tuple-element.5 = f32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=1 - %get-tuple-element.6 = u32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=2 - %get-tuple-element.7 = s32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=3 - %copy.4 = s32[1]{0:T(128)} copy(s32[1]{0:T(128)} %get-tuple-element.4) - %copy.5 = f32[1]{0:T(128)} copy(f32[1]{0:T(128)} %get-tuple-element.5) - %copy.6 = u32[1]{0:T(128)} copy(u32[1]{0:T(128)} %get-tuple-element.6) - %copy.7 = s32[1]{0:T(128)} copy(s32[1]{0:T(128)} %get-tuple-element.7) - ROOT %tuple.2 = (s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) tuple(s32[1]{0:T(128)} %copy.4, f32[1]{0:T(128)} %copy.5, u32[1]{0:T(128)} %copy.6, s32[1]{0:T(128)} %copy.7) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, BitcastMultiUse) { - // When there is a pattern where a bitcast has multiple uses (negate0 and add) - // and one is in the default memory and the other is in alternate memory, they - // both need their own bitcast. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - HloInstruction* p0 = builder.AddInstruction( - HloInstruction::CreateParameter(0, param_shape, "p1")); - HloInstruction* bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(shape, p0)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, bitcast)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, bitcast, negate4)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, bitcast, negate0, negate1, negate2, - negate3, negate4, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - EXPECT_THAT(negate0->operand(0), op::ShapeWithLayout(shape)); - EXPECT_THAT(add->operand(0), op::ShapeWithLayout(shape_in_alternate_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, BitcastMultiUseTuple) { - // Same as BitcastMultUse but the second use is a tuple. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - - auto module = CreateNewVerifiedModule(); - HloComputation::Builder fusion_builder("fusion"); - HloInstruction* fusion_param = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* fusion_element0 = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion_param, 0)); - HloInstruction* fusion_element1 = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion_param, 1)); - fusion_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, fusion_element0, fusion_element1)); - HloComputation* fusion_computation = - module->AddEmbeddedComputation(fusion_builder.Build()); - - HloInstruction* p0 = builder.AddInstruction( - HloInstruction::CreateParameter(0, param_shape, "p1")); - HloInstruction* bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(shape, p0)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, bitcast)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({bitcast, negate4})); - HloInstruction* fusion = builder.AddInstruction(HloInstruction::CreateFusion( - shape, HloInstruction::FusionKind::kCustom, {tuple}, fusion_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, bitcast, negate0, negate1, negate2, - negate3, negate4, tuple, fusion}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - EXPECT_THAT(negate0->operand(0), op::ShapeWithLayout(shape)); - EXPECT_THAT(fusion->operand(0)->operand(0), - op::ShapeWithLayout(shape_in_alternate_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, BitcastScheduleBug) { - // Bitcasts can force asynchronous copies to be scheduled too early, possibly - // leading to memory corruption. - // Bug: - // p0------------------>neg-->neg-->neg ... -->neg-->neg-->neg->add - // / - // p1->cs->cd->bitcast-----------------------------------------+ - // - // Expected: - // p0-->neg-->neg-->neg ... -->neg-->neg-->neg------------->add - // / - // p1--------------------->cs----------------->cd->bitcast-+ - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape param_shape = ShapeUtil::MakeShape(F32, {6}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateParameter(1, param_shape, "p1")); - HloInstruction* bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(shape, p1)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* negate8 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate7)); - HloInstruction* negate9 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate8)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, bitcast, negate9)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, {p0, p1, bitcast, negate0, negate1, negate2, negate3, - negate4, negate5, negate6, negate7, negate8, negate9, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/4); - - EXPECT_EQ(add->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - const auto& instructions = - module->schedule().sequence(module->entry_computation()).instructions(); - for (int i = 0; i < instructions.size(); ++i) { - // Expect that there is a negate before and after the CopyStart and there is - // a negate before CopyDone. - if (instructions.at(i)->opcode() == HloOpcode::kCopyStart) { - EXPECT_EQ(instructions.at(i - 1)->opcode(), HloOpcode::kNegate); - EXPECT_EQ(instructions.at(i + 1)->opcode(), HloOpcode::kNegate); - } else if (instructions.at(i)->opcode() == HloOpcode::kCopyDone) { - EXPECT_EQ(instructions.at(i - 1)->opcode(), HloOpcode::kNegate); - } - } -} - -TEST_F(MemorySpaceAssignmentTest, AddDependency) { - // Make sure add-dependency is not optimized away. - absl::string_view hlo_string = R"hlo( - HloModule AddDependency, is_scheduled=true - - ENTRY %AddDependency (p: f32[3]) -> f32[3] { - %p = f32[3]{0} parameter(0) - %neg0 = f32[3]{0} negate(f32[3]{0} %p) - %neg1 = f32[3]{0} negate(f32[3]{0} %neg0) - %neg2 = f32[3]{0} negate(f32[3]{0} %neg1) - %neg3 = f32[3]{0} negate(f32[3]{0} %neg2) - %neg4 = f32[3]{0} negate(f32[3]{0} %neg3) - %neg5 = f32[3]{0} negate(f32[3]{0} %neg4) - %neg6 = f32[3]{0} negate(f32[3]{0} %neg5) - %token0 = token[] after-all() - %add_dep = f32[3]{0} add-dependency(f32[3]{0} %p, token[] %token0) - ROOT %add = f32[3]{0} add(f32[3]{0} %add_dep, f32[3]{0} %neg6) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Add(op::AddDependency(), op::Negate())); -} - -TEST_F(MemorySpaceAssignmentTest, WhileAllocationBug) { - // This test is carefully crafted to include two multiply ops sized [4,3] in a - // while body. For testing purposes, we have provided a BufferIntervalCompare - // such that first multiply, then tanh, then other HloValues will be - // allocated. The memory is sized just enough to fit two [4,3] buffers. - // Because the multiplies in the while body are going to be allocated in the - // alternate memory first, the tanh that is fed inside the while loop should - // not be placed in the alternate memory. Otherwise, we will corrupt memory. - absl::string_view hlo_string = R"hlo( - HloModule WhileAllocationBug, is_scheduled=true - - %WhileBody (body_param: (f32[4,3], f32[])) -> (f32[4,3], f32[]) { - %body_param = (f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[4,3]{1,0}, f32[]) %body_param), index=1 - %get-tuple-element.2 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[]) %body_param), index=0 - %constant.1 = f32[] constant(1) - %add = f32[] add(f32[] %get-tuple-element.1, f32[] %constant.1) - %constant.2 = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 1, 2, 3 }, { 4, 5, 6 } }) - %multiply = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %get-tuple-element.2, f32[4,3]{1,0} %get-tuple-element.2) - %multiply2 = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %multiply, f32[4,3]{1,0} %multiply) - %add.1 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.2, f32[4,3]{1,0} %constant.2) - %add.2 = f32[4,3]{1,0} add(f32[4,3]{1,0} %add.1, f32[4,3]{1,0} %multiply2) - ROOT %tuple = (f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %add.2, f32[] %add) - } - - %WhileCond (cond_param: (f32[4,3], f32[])) -> pred[] { - %cond_param = (f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[4,3]{1,0}, f32[]) %cond_param), index=1 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - ENTRY %Entry (param_iter: f32[4,3], param_data: f32[], p2: f32[4,3]) -> f32[4,3] { - %param_data = f32[] parameter(1) - %param_iter = f32[4,3]{1,0} parameter(0) - %p2 = f32[4,3]{1,0} parameter(2) - %tanh = f32[4,3]{1,0} tanh(f32[4,3]{1,0} %param_iter) - %neg0 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %p2) - %neg1 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg0) - %neg2 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg1) - %neg3 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg2) - %neg4 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg3) - %neg5 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg4) - %neg6 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg5) - %add.4 = f32[4,3]{1,0} add(f32[4,3]{1,0} %neg6, f32[4,3]{1,0} %tanh) - %tuple.1 = (f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %tanh, f32[] %param_data) - %while = (f32[4,3]{1,0}, f32[]) while((f32[4,3]{1,0}, f32[]) %tuple.1), condition=%WhileCond, body=%WhileBody - %get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[]) %while), index=0 - ROOT %add.3 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.3, f32[4,3]{1,0} %add.4) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - bool a_is_mul = - a.buffer->defining_instruction()->opcode() == HloOpcode::kMultiply; - bool b_is_mul = - b.buffer->defining_instruction()->opcode() == HloOpcode::kMultiply; - if (a_is_mul && !b_is_mul) { - return true; - } - if (!a_is_mul && b_is_mul) { - return false; - } - bool a_is_tanh = - a.buffer->defining_instruction()->opcode() == HloOpcode::kTanh; - bool b_is_tanh = - b.buffer->defining_instruction()->opcode() == HloOpcode::kTanh; - if (a_is_tanh && !b_is_tanh) { - return true; - } - if (!a_is_tanh && b_is_tanh) { - return false; - } - return a.buffer->id() < b.buffer->id(); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - buffer_interval_compare, &prefetch_interval_picker); - - for (const HloInstruction* instruction : - module->entry_computation()->instructions()) { - if (instruction->opcode() == HloOpcode::kWhile) { - const Shape& while_subshape = - ShapeUtil::GetSubshape(instruction->shape(), {0}); - // We expect shape {0} to either be in default memory for the entire while - // loop or there has to be an eviction within the while loop. - if (while_subshape.layout().memory_space() == kAlternateMemorySpace) { - const HloInstruction* body_param = - instruction->while_body()->parameter_instruction(0); - const HloInstruction* gte = nullptr; - for (const HloInstruction* user : body_param->users()) { - if (user->opcode() == HloOpcode::kGetTupleElement && - user->tuple_index() == 0) { - gte = user; - break; - } - } - EXPECT_NE(gte, nullptr); - const HloInstruction* copy_start = nullptr; - for (const HloInstruction* user : gte->users()) { - if (user->opcode() == HloOpcode::kCopyStart) { - copy_start = user; - break; - } - } - EXPECT_NE(copy_start, nullptr); - const Shape& copy_start_subshape = - ShapeUtil::GetSubshape(copy_start->shape(), {0}); - - EXPECT_NE(copy_start_subshape.layout().memory_space(), - kAlternateMemorySpace); - } - } - } -} - -TEST_F(MemorySpaceAssignmentTest, ConsecutiveWhileLoops) { - absl::string_view hlo_string = R"hlo( - HloModule WhileAllocationBug, is_scheduled=true - - %WhileBody (body_param: (f32[4,3], f32[4,3], f32[])) -> (f32[4,3], f32[4,3], f32[]) { - %body_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=2 - %get-tuple-element.2 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=0 - %get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=1 - %constant.1 = f32[] constant(1) - %add = f32[] add(f32[] %get-tuple-element.1, f32[] %constant.1) - %constant.2 = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 1, 2, 3 }, { 4, 5, 6 } }) - %multiply = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %get-tuple-element.2, f32[4,3]{1,0} %get-tuple-element.3) - %multiply2 = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %multiply, f32[4,3]{1,0} %multiply) - %add.1 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.2, f32[4,3]{1,0} %constant.2) - %add.2 = f32[4,3]{1,0} add(f32[4,3]{1,0} %add.1, f32[4,3]{1,0} %multiply2) - ROOT %tuple = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %add.2, f32[4,3]{1,0} %get-tuple-element.3, f32[] %add) - } - - %WhileCond (cond_param: (f32[4,3], f32[4,3], f32[])) -> pred[] { - %cond_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %cond_param), index=2 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - %WhileBody2 (body_param: (f32[4,3], f32[4,3], f32[])) -> (f32[4,3], f32[4,3], f32[]) { - %body_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=2 - %get-tuple-element.2 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=0 - %get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=1 - %constant.1 = f32[] constant(1) - %add = f32[] add(f32[] %get-tuple-element.1, f32[] %constant.1) - %constant.2 = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 1, 2, 3 }, { 4, 5, 6 } }) - %multiply = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %get-tuple-element.2, f32[4,3]{1,0} %get-tuple-element.3) - %multiply2 = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %multiply, f32[4,3]{1,0} %multiply) - %add.1 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.2, f32[4,3]{1,0} %constant.2) - %add.2 = f32[4,3]{1,0} add(f32[4,3]{1,0} %add.1, f32[4,3]{1,0} %multiply2) - ROOT %tuple = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %add.2, f32[4,3]{1,0} %get-tuple-element.3, f32[] %add) - } - - %WhileCond2 (cond_param: (f32[4,3], f32[4,3], f32[])) -> pred[] { - %cond_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %cond_param), index=2 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - ENTRY %Entry (param_data: f32[4,3], param_iter: f32[], p2: f32[4,3]) -> f32[4,3] { - %param_iter = f32[] parameter(1) - %param_data = f32[4,3]{1,0} parameter(0) - %p2 = f32[4,3]{1,0} parameter(2) - %neg0 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %p2) - %neg1 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg0) - %neg2 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg1) - %neg3 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg2) - %neg4 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg3) - %neg5 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg4) - %neg6 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg5) - %add.4 = f32[4,3]{1,0} add(f32[4,3]{1,0} %neg6, f32[4,3]{1,0} %p2) - %tuple.1 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} add.4, f32[4,3]{1,0} param_data, f32[] %param_iter) - %while = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) while((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %tuple.1), condition=%WhileCond, body=%WhileBody - %get-tuple-element.4 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while), index=0 - %add.3 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.4, f32[4,3]{1,0} %add.4) - %get-tuple-element.5 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while), index=1 - %tuple.2 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} add.3, f32[4,3]{1,0} get-tuple-element.5, f32[] %param_iter) - %while.1 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) while((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %tuple.2), condition=%WhileCond2, body=%WhileBody2 - %get-tuple-element.6 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while.1), index=0 - ROOT %add.5 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.6, f32[4,3]{1,0} %add.3) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, WhileLiveRangeBug) { - // Tests against while live ranges being incorrect and the verifier - // complaining about a conflict. - absl::string_view hlo_string = R"hlo( - HloModule WhileAllocationBug, is_scheduled=true - - %WhileBody (body_param: (f32[4,3], f32[4,3], f32[])) -> (f32[4,3], f32[4,3], f32[]) { - %body_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=2 - %get-tuple-element.2 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=0 - %get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=1 - %neg10 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %get-tuple-element.2) - %neg11 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg10) - %neg12 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg11) - %neg13 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg12) - %neg14 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg13) - %neg15 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg14) - %neg16 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg15) - %neg17 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg16) - %neg18 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg17) - %neg19 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg18) - %neg20 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg19) - %constant.1 = f32[] constant(1) - %add = f32[] add(f32[] %get-tuple-element.1, f32[] %constant.1) - %constant.2 = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 1, 2, 3 }, { 4, 5, 6 } }) - %multiply = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %neg20, f32[4,3]{1,0} %neg20) - %multiply2 = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %multiply, f32[4,3]{1,0} %multiply) - %add.1 = f32[4,3]{1,0} add(f32[4,3]{1,0} get-tuple-element.3, f32[4,3]{1,0} %constant.2) - %add.2 = f32[4,3]{1,0} add(f32[4,3]{1,0} %add.1, f32[4,3]{1,0} %multiply2) - ROOT %tuple = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %add.2, f32[4,3]{1,0} %get-tuple-element.3, f32[] %add) - } - - %WhileCond (cond_param: (f32[4,3], f32[4,3], f32[])) -> pred[] { - %cond_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %cond_param), index=2 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - ENTRY %Entry (param_data: f32[4,3], param_iter: f32[], p2: f32[4,3]) -> f32[4,3] { - %param_iter = f32[] parameter(1) - %param_data = f32[4,3]{1,0} parameter(0) - %p2 = f32[4,3]{1,0} parameter(2) - %neg0 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %p2) - %neg1 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg0) - %neg2 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg1) - %neg3 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg2) - %neg4 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg3) - %neg5 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg4) - %neg6 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg5) - %add.4 = f32[4,3]{1,0} add(f32[4,3]{1,0} %neg6, f32[4,3]{1,0} %p2) - %tuple.1 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} add.4, f32[4,3]{1,0} param_data, f32[] %param_iter) - %while = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) while((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %tuple.1), condition=%WhileCond, body=%WhileBody - %get-tuple-element.4 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while), index=0 - %get-tuple-element.5 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while), index=1 - %add.3 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.4, f32[4,3]{1,0} %add.4) - ROOT %add.5 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.5, f32[4,3]{1,0} %add.3) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, ConsecutiveWhileLoopsOneBuffer) { - // Tests against a bug when there are consecutive while loops with one buffer - // (the value doesn't change in the buffer), the parameter can be colored in - // the alternate memory space. - absl::string_view hlo_string = R"hlo( - HloModule WhileAllocationBug, is_scheduled=true - - %WhileBody (body_param: (f32[4,3], f32[4,3], f32[])) -> (f32[4,3], f32[4,3], f32[]) { - %body_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=2 - %get-tuple-element.2 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=0 - %get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=1 - %neg10 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %get-tuple-element.2) - %neg11 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg10) - %neg12 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg11) - %neg13 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg12) - %neg14 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg13) - %neg15 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg14) - %neg16 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg15) - %neg17 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg16) - %neg18 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg17) - %neg19 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg18) - %neg20 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg19) - %constant.1 = f32[] constant(1) - %add = f32[] add(f32[] %get-tuple-element.1, f32[] %constant.1) - %constant.2 = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 1, 2, 3 }, { 4, 5, 6 } }) - %multiply = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %neg20, f32[4,3]{1,0} %neg20) - %multiply2 = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %multiply, f32[4,3]{1,0} %multiply) - %add.1 = f32[4,3]{1,0} add(f32[4,3]{1,0} get-tuple-element.3, f32[4,3]{1,0} %constant.2) - %add.2 = f32[4,3]{1,0} add(f32[4,3]{1,0} %add.1, f32[4,3]{1,0} %multiply2) - ROOT %tuple = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %add.2, f32[4,3]{1,0} %get-tuple-element.3, f32[] %add) - } - - %WhileCond (cond_param: (f32[4,3], f32[4,3], f32[])) -> pred[] { - %cond_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %cond_param), index=2 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - %WhileBody2 (body_param: (f32[4,3], f32[4,3], f32[])) -> (f32[4,3], f32[4,3], f32[]) { - %body_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=2 - %get-tuple-element.2 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=0 - %get-tuple-element.3 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %body_param), index=1 - %neg10 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %get-tuple-element.2) - %neg11 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg10) - %neg12 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg11) - %neg13 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg12) - %neg14 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg13) - %neg15 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg14) - %neg16 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg15) - %neg17 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg16) - %neg18 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg17) - %neg19 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg18) - %neg20 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg19) - %constant.1 = f32[] constant(1) - %add = f32[] add(f32[] %get-tuple-element.1, f32[] %constant.1) - %constant.2 = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 1, 2, 3 }, { 4, 5, 6 } }) - %multiply = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %neg20, f32[4,3]{1,0} %neg20) - %multiply2 = f32[4,3]{1,0} multiply(f32[4,3]{1,0} %multiply, f32[4,3]{1,0} %multiply) - %add.1 = f32[4,3]{1,0} add(f32[4,3]{1,0} get-tuple-element.3, f32[4,3]{1,0} %constant.2) - %add.2 = f32[4,3]{1,0} add(f32[4,3]{1,0} %add.1, f32[4,3]{1,0} %multiply2) - ROOT %tuple = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} %add.2, f32[4,3]{1,0} %get-tuple-element.3, f32[] %add) - } - - %WhileCond2 (cond_param: (f32[4,3], f32[4,3], f32[])) -> pred[] { - %cond_param = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %cond_param), index=2 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - ENTRY %Entry (param_data: f32[4,3], param_iter: f32[], p2: f32[4,3]) -> f32[4,3] { - %param_iter = f32[] parameter(1) - %param_data = f32[4,3]{1,0} parameter(0) - %p2 = f32[4,3]{1,0} parameter(2) - %neg0 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %p2) - %neg1 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg0) - %neg2 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg1) - %neg3 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg2) - %neg4 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg3) - %neg5 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg4) - %neg6 = f32[4,3]{1,0} negate(f32[4,3]{1,0} %neg5) - %add.4 = f32[4,3]{1,0} add(f32[4,3]{1,0} %neg6, f32[4,3]{1,0} %p2) - %tuple.1 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} add.4, f32[4,3]{1,0} param_data, f32[] %param_iter) - %while = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) while((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %tuple.1), condition=%WhileCond, body=%WhileBody - %get-tuple-element.4 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while), index=0 - %add.3 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.4, f32[4,3]{1,0} %add.4) - %tuple.2 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) tuple(f32[4,3]{1,0} add.3, f32[4,3]{1,0} param_data, f32[] %param_iter) - %while.1 = (f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) while((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %tuple.2), condition=%WhileCond2, body=%WhileBody2 - %get-tuple-element.5 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while.1), index=0 - %get-tuple-element.6 = f32[4,3]{1,0} get-tuple-element((f32[4,3]{1,0}, f32[4,3]{1,0}, f32[]) %while.1), index=1 - ROOT %add.5 = f32[4,3]{1,0} add(f32[4,3]{1,0} %get-tuple-element.5, f32[4,3]{1,0} %get-tuple-element.6) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, WhileCondAliasBug) { - // While loop is the root of the entry computation. We should ensure the - // output of the entry computation remains to be in default memory space. - // Test from //third_party/tensorflow/compiler/xla/tests:while_test - // WhileTest.WhileWithPrngScalarResult. - absl::string_view hlo_string = R"hlo( - HloModule WhileWithPrngScalarResult.18, is_scheduled=true - - %fused_computation (param_0.1: s32[6], param_1.3: s32[1], param_2.3: s32[5]) -> s32[6] { - %param_1.3 = s32[1]{0:T(128)} parameter(1) - %constant.2 = s32[]{:T(128)} constant(-2147483648) - %pad.2 = s32[6]{0:T(128)} pad(s32[1]{0:T(128)} %param_1.3, s32[]{:T(128)} %constant.2), padding=0_5 - %param_2.3 = s32[5]{0:T(128)} parameter(2) - %pad.3 = s32[6]{0:T(128)} pad(s32[5]{0:T(128)} %param_2.3, s32[]{:T(128)} %constant.2), padding=1_0 - %maximum.1 = s32[6]{0:T(128)} maximum(s32[6]{0:T(128)} %pad.2, s32[6]{0:T(128)} %pad.3) - %param_0.1 = s32[6]{0:T(128)} parameter(0) - ROOT %add.0 = s32[6]{0:T(128)} add(s32[6]{0:T(128)} %maximum.1, s32[6]{0:T(128)} %param_0.1) - } - - %body.3 (prev.4: s32[6]) -> s32[6] { - %constant.7 = s32[]{:T(128)} constant(100) - %constant.6 = s32[]{:T(128)} constant(0) - %constant.5 = s32[1]{0:T(128)} constant({1}) - %prev.4 = s32[6]{0:T(128)} parameter(0) - %rng.8 = s32[5]{0:T(128)} rng(s32[]{:T(128)} %constant.6, s32[]{:T(128)} %constant.7), distribution=rng_uniform - %neg = s32[1]{0:T(128)} negate(s32[1]{0:T(128)} %constant.5) - ROOT %fusion = s32[6]{0:T(128)} fusion(s32[6]{0:T(128)} %prev.4, s32[1]{0:T(128)} %neg, s32[5]{0:T(128)} %rng.8), kind=kLoop, calls=%fused_computation - } - - %WhileWithPrngScalarResult.11 (prev.12: s32[6]) -> pred[] { - %constant.15 = s32[]{:T(128)} constant(1) - %prev.12 = s32[6]{0:T(128)} parameter(0) - %bitcast.1 = s32[1]{0:T(128)} bitcast(s32[6]{0:T(128)} %prev.12) - %bitcast = s32[]{:T(128)} bitcast(s32[1]{0:T(128)} %bitcast.1) - ROOT %compare.16 = pred[]{:T(128)} compare(s32[]{:T(128)} %constant.15, s32[]{:T(128)} %bitcast), direction=GT - } - - ENTRY %WhileWithPrngScalarResult.18 () -> s32[6] { - %constant.1 = s32[]{:T(128)} constant(0) - %broadcast.2 = s32[6]{0:T(128)} broadcast(s32[]{:T(128)} %constant.1), dimensions={} - ROOT %while.17 = s32[6]{0:T(128)} while(s32[6]{0:T(128)} %broadcast.2), condition=%WhileWithPrngScalarResult.11, body=%body.3 - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, WhileInPlaceBuffer) { - // Ensure that a dynamic update slice within a while loop is able to get an - // alternate memory allocation. - absl::string_view hlo_string = R"hlo( - HloModule Module, is_scheduled=true - - fused_computation { - param0 = f32[2,3] parameter(0) - constant.1 = f32[] constant(0) - broadcast = f32[2,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - ROOT dynamic-update-slice.5 = f32[2,3] dynamic-update-slice(param0, broadcast, constant.3, constant.3) - } - - %WhileBody (body_param: (f32[2,3], f32[2,3], f32[])) -> (f32[2,3], f32[2,3], f32[]) { - %body_param = (f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) parameter(0) - %get-tuple-element.1 = f32[] get-tuple-element((f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) %body_param), index=2 - %get-tuple-element.2 = f32[2,3]{1,0} get-tuple-element((f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) %body_param), index=0 - %get-tuple-element.3 = f32[2,3]{1,0} get-tuple-element((f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) %body_param), index=1 - %fusion = f32[2,3]{1,0} fusion(get-tuple-element.3), kind=kLoop, calls=fused_computation - %multiply = f32[2,3]{1,0} multiply(f32[2,3]{1,0} %get-tuple-element.2, f32[2,3]{1,0} %fusion) - ROOT %tuple = (f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) tuple(f32[2,3]{1,0} %multiply, f32[2,3]{1,0} %fusion, f32[] %get-tuple-element.1) - } - - %WhileCond (cond_param: (f32[2,3], f32[2,3], f32[])) -> pred[] { - %cond_param = (f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) parameter(0) - %get-tuple-element = f32[] get-tuple-element((f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) %cond_param), index=2 - %constant = f32[] constant(50) - ROOT %compare = pred[] compare(f32[] %get-tuple-element, f32[] %constant), direction=LT - } - - ENTRY %Entry (param_data: f32[2,3], param_iter: f32[], p2: f32[2,3]) -> f32[2,3] { - %param_iter = f32[] parameter(1) - %param_data = f32[2,3]{1,0} parameter(0) - %p2 = f32[2,3]{1,0} parameter(2) - %copy1 = f32[2,3]{1,0} copy(param_data) - %copy2 = f32[2,3]{1,0} copy(p2) - %tuple.1 = (f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) tuple(f32[2,3]{1,0} copy1, f32[2,3]{1,0} copy2, f32[] %param_iter) - %while = (f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) while((f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) %tuple.1), condition=%WhileCond, body=%WhileBody - %get-tuple-element.4 = f32[2,3]{1,0} get-tuple-element((f32[2,3]{1,0}, f32[2,3]{1,0}, f32[]) %while), index=0 - ROOT %copy3 = f32[2,3]{1,0} copy(get-tuple-element.4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - const HloInstruction* while_op = - module->entry_computation()->GetInstructionWithName("while"); - EXPECT_EQ( - ShapeUtil::GetSubshape(while_op->shape(), {1}).layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, WhileSharedBufferVerificationBug) { - // Tests a spurious verification failure when a while has the same value - // passed in twice (copy0) and that value is evicted within the while loop. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=3 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = f32[3]{0} get-tuple-element(p0), index=2 - gte3 = pred[] get-tuple-element(p0), index=3 - add = f32[3]{0} add(gte0, gte0) - negate0 = f32[3]{0} negate(add) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - negate15 = f32[3]{0} negate(negate14) - negate16 = f32[3]{0} negate(negate15) - ROOT tuple = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, gte0, negate16, gte3) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy0 = f32[3]{0} copy(p0) - copy1 = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) tuple(copy0, copy0, copy1, p1) - while = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - ROOT gte = f32[3]{0} get-tuple-element(while), index=2 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, b228599972) { - absl::string_view hlo_string = R"hlo( -HloModule entry, is_scheduled=true - -fused_computation { - %p0 = f32[2,3]{1,0} parameter(0) - %result0 = f32[2,3]{1,0} copy(%p0) - %result1 = f32[2,3]{1,0} copy(%p0) - ROOT tuple = (f32[2,3]{1,0}, f32[2,3]{1,0}) tuple(%result0, %result1) -} - -ENTRY entry { - %p0 = f32[2,3]{1,0} parameter(0) - %p1 = f32[2,3]{1,0} parameter(1) - %unused = (f32[2,3]{1,0}, f32[2,3]{1,0}) fusion(%p0), kind=kLoop, calls=%fused_computation - %unused.0 = f32[2,3]{1,0} get-tuple-element(%unused), index=0 - %unused.1 = f32[2,3]{1,0} get-tuple-element(%unused), index=1 - %negate.0 = f32[2,3]{1,0} negate(f32[2,3]{1,0} %unused.0) - %negate.1 = f32[2,3]{1,0} negate(f32[2,3]{1,0} %unused.1) - - ROOT %result = f32[2,3]{1,0} negate(%p1) -} - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, b172243149) { - // Tests for the failure in b/172243149, where if we skip processing - // non-copy allocations that are in default memory can actually cause - // failures. In this case, the problem tensor is copy0, where it is fed to - // both negate, while, and add0. The copy0->negate dependency can be allocated - // in the alternate memory. Then the algorithm attempts to place the - // copy0->while edge in the alternate memory, but since this value isn't used - // in the while loop, it won't get an alternate memory allocation. Finally for - // the copy0->add0 edge, the algorithm will actually replace it with - // while{0}->add0, since this is equivalent and while is defined later than - // copy0. However, if we actually skip processing this while{0}->add0 - // allocation, we won't replace this edge, and will end up with the - // copy0->add0 edge, which illegally extends the lifetime of the alternate - // memory buffer in copy0. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=3 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = f32[3]{0} get-tuple-element(p0), index=2 - gte3 = pred[] get-tuple-element(p0), index=3 - add = f32[3]{0} add(gte1, gte2) - negate0 = f32[3]{0} negate(add) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - negate15 = f32[3]{0} negate(negate14) - negate16 = f32[3]{0} negate(negate15) - ROOT tuple = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, add, negate16, gte3) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy0 = f32[3]{0} copy(p0) - copy1 = f32[3]{0} copy(p0) - copy2 = f32[3]{0} copy(p0) - negate = f32[3]{0} negate(copy0) - tuple = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) tuple(copy0, copy1, copy2, p1) - while = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - gte = f32[3]{0} get-tuple-element(while), index=2 - add0 = f32[3]{0} add(negate, copy0) - ROOT add1 = f32[3]{0} add(add0, gte) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, ControlPredecessorsBug) { - // Having control_predecessors on an HLO was preventing us from DCEing an op - // that doesn't have any users (tuple.1). The scheduler assumes the graph is - // fully DCEed, which causes some instructions not to be scheduled. - absl::string_view hlo_string = R"hlo( - HloModule sort.16, is_scheduled=true - - ENTRY %sort.16 (param.0.1: s32[1], param.1.2: f32[1], param.2.3: u32[1], param.3.4: s32[1]) -> (s32[1], f32[1], u32[1], s32[1]) { - %param.3.4 = s32[1]{0:T(128)} parameter(3) - %param.2.3 = u32[1]{0:T(128)} parameter(2) - %param.1.2 = f32[1]{0:T(128)} parameter(1) - %param.0.1 = s32[1]{0:T(128)} parameter(0) - %tuple.1 = (s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) tuple(s32[1]{0:T(128)} %param.0.1, f32[1]{0:T(128)} %param.1.2, u32[1]{0:T(128)} %param.2.3, s32[1]{0:T(128)} %param.3.4), control-predecessors={%param.0.1} - %get-tuple-element.4 = s32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=0 - %get-tuple-element.5 = f32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=1 - %get-tuple-element.6 = u32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=2 - %get-tuple-element.7 = s32[1]{0:T(128)} get-tuple-element((s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) %tuple.1), index=3 - %copy.4 = s32[1]{0:T(128)} copy(s32[1]{0:T(128)} %get-tuple-element.4) - %copy.5 = f32[1]{0:T(128)} copy(f32[1]{0:T(128)} %get-tuple-element.5) - %copy.6 = u32[1]{0:T(128)} copy(u32[1]{0:T(128)} %get-tuple-element.6) - %copy.7 = s32[1]{0:T(128)} copy(s32[1]{0:T(128)} %get-tuple-element.7) - ROOT %tuple.2 = (s32[1]{0:T(128)}, f32[1]{0:T(128)}, u32[1]{0:T(128)}, s32[1]{0:T(128)}) tuple(s32[1]{0:T(128)} %copy.4, f32[1]{0:T(128)} %copy.5, u32[1]{0:T(128)} %copy.6, s32[1]{0:T(128)} %copy.7) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalShouldBeAllocatedInAlternateMem) { - // Checks if simple conditionals get alternate memory allocations. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg1 = f32[3]{0} negate(gte) - } - - false_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg2 = f32[3]{0} negate(gte) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}) tuple(copy) - ROOT conditional = f32[3]{0} conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Check that copy and gtes got alternate memory allocations. - auto copy = - module->GetComputationWithName("entry")->GetInstructionWithName("copy"); - EXPECT_EQ(copy->shape().layout().memory_space(), kAlternateMemorySpace); - auto neg1 = module->GetComputationWithName("true_computation") - ->GetInstructionWithName("neg1"); - auto neg1_operand = neg1->operand(0); - EXPECT_EQ(neg1_operand->shape().layout().memory_space(), - kAlternateMemorySpace); - auto neg2 = module->GetComputationWithName("false_computation") - ->GetInstructionWithName("neg2"); - auto neg2_operand = neg2->operand(0); - EXPECT_EQ(neg2_operand->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalAvoidsUnnecessaryPrefetch) { - // Checks if we avoid unnecessary allocation in alternate memory if the input - // won't be used in the computation for a long time. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]{0}, f32[3]{0}) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - neg0 = f32[3]{0} negate(gte0) - neg1 = f32[3]{0} negate(neg0) - neg2 = f32[3]{0} negate(neg1) - neg3 = f32[3]{0} negate(neg2) - neg4 = f32[3]{0} negate(neg3) - neg5 = f32[3]{0} negate(neg4) - neg6 = f32[3]{0} negate(neg5) - neg7 = f32[3]{0} negate(neg6) - neg8 = f32[3]{0} negate(neg7) - neg9 = f32[3]{0} negate(neg8) - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - ROOT add = f32[3]{0} add(neg9, gte1) - } - - false_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg = f32[3]{0} negate(gte) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy0 = f32[3]{0} copy(p0) - copy1 = f32[3]{0} copy(p0) - tuple0 = (f32[3]{0}, f32[3]{0}) tuple(copy0, copy1) - tuple1 = (f32[3]{0}) tuple(copy0) - ROOT conditional = f32[3]{0} conditional(p1, tuple0, tuple1), true_computation=true_computation, false_computation=false_computation - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Check that copy1 doesn't get unnecessarily allocated in alternate mem - // (due to long negate chain in true_computation) but is prefetched before - // add. - auto copy0 = - module->GetComputationWithName("entry")->GetInstructionWithName("copy0"); - EXPECT_EQ(copy0->shape().layout().memory_space(), kAlternateMemorySpace); - auto copy1 = - module->GetComputationWithName("entry")->GetInstructionWithName("copy1"); - EXPECT_EQ(copy1->shape().layout().memory_space(), kDefaultMemorySpace); - auto add = module->GetComputationWithName("true_computation") - ->GetInstructionWithName("add"); - auto add_operand = add->operand(1); - EXPECT_EQ(add_operand->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalMultiUse) { - // Make sure there is an evict when there is a conditional use followed by - // another use. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]{0}, f32[3]{0}) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - add0 = f32[3]{0} add(gte0, gte1) - neg0 = f32[3]{0} negate(add0) - neg1 = f32[3]{0} negate(neg0) - neg2 = f32[3]{0} negate(neg1) - neg3 = f32[3]{0} negate(neg2) - neg4 = f32[3]{0} negate(neg3) - neg5 = f32[3]{0} negate(neg4) - neg6 = f32[3]{0} negate(neg5) - neg7 = f32[3]{0} negate(neg6) - neg8 = f32[3]{0} negate(neg7) - ROOT neg9 = f32[3]{0} negate(neg8) - } - - false_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg = f32[3]{0} negate(gte) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy0 = f32[3]{0} copy(p0) - copy1 = f32[3]{0} copy(p0) - tuple0 = (f32[3]{0}, f32[3]{0}) tuple(copy0, copy1) - tuple1 = (f32[3]{0}) tuple(copy0) - conditional = f32[3]{0} conditional(p1, tuple0, tuple1), true_computation=true_computation, false_computation=false_computation - ROOT add1 = f32[3]{0} add(copy1, conditional) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Make sure the add0 and add1 uses of copy1 are in alternate memory. - // Make sure there are no evictions for copy1. - auto copy1 = - module->GetComputationWithName("entry")->GetInstructionWithName("copy1"); - EXPECT_EQ(copy1->shape().layout().memory_space(), kAlternateMemorySpace); - auto add0 = module->GetComputationWithName("true_computation") - ->GetInstructionWithName("add0"); - auto add0_operand = add0->operand(1); - EXPECT_EQ(add0_operand->shape().layout().memory_space(), - kAlternateMemorySpace); - auto add1 = - module->GetComputationWithName("entry")->GetInstructionWithName("add1"); - auto add1_operand = add1->operand(0); - EXPECT_EQ(add1_operand->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(add1_operand->opcode(), HloOpcode::kCopy); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalMultiUseInWhile) { - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg1 = f32[3]{0} negate(gte) - } - - false_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg2 = f32[3]{0} negate(gte) - } - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = pred[] get-tuple-element(p0), index=2 - cond_tuple = (f32[3]{0}) tuple(gte0) - conditional = f32[3]{0} conditional(gte2, cond_tuple, cond_tuple), true_computation=true_computation, false_computation=false_computation - add = f32[3]{0} add(conditional, gte1) - neg0 = f32[3]{0} negate(add) - neg1 = f32[3]{0} negate(neg0) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, neg1, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy0 = f32[3]{0} copy(p0) - copy1 = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy0, copy1, p1) - while = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - ROOT gte = f32[3]{0} get-tuple-element(while), index=1 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - // Make sure copy0/while{0}/cond_tuple{0} gets alternate memory allocation. - // Make sure there are no unnecessary evictions. - auto copy0 = - module->GetComputationWithName("entry")->GetInstructionWithName("copy0"); - EXPECT_EQ(copy0->shape().layout().memory_space(), kAlternateMemorySpace); - auto conditional = module->GetComputationWithName("while_body") - ->GetInstructionWithName("conditional"); - auto conditional_operand = conditional->operand(1); - EXPECT_EQ(ShapeUtil::GetSubshape(conditional_operand->shape(), {0}) - .layout() - .memory_space(), - kAlternateMemorySpace); - auto while_root = - module->GetComputationWithName("while_body")->root_instruction(); - auto while_root_operand = while_root->operand(0); - // Verify that the while body root reads from the parameter instead of an - // evicted copy. - EXPECT_THAT(while_root_operand, op::GetTupleElement(op::Parameter(0))); -} - -TEST_F(MemorySpaceAssignmentTest, NestedConditional) { - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation2 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg1 = f32[3]{0} negate(gte) - } - - false_computation2 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg2 = f32[3]{0} negate(gte) - } - - true_computation1 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - slice = f32[1]{0} slice(gte), slice={[0:1]} - bitcast = f32[] bitcast(slice) - constant = f32[] constant(0.0) - compare = pred[] compare(bitcast, constant), direction=GT - ROOT conditional = f32[3]{0} conditional(compare, p0, p0), true_computation=true_computation2, false_computation=false_computation2 - } - - false_computation1 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg3 = f32[3]{0} negate(gte) - } - - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}) tuple(copy) - ROOT conditional = f32[3]{0} conditional(p1, tuple, tuple), true_computation=true_computation1, false_computation=false_computation1 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Make sure alternate memory allocation gets propagated into both levels of - // conditional. - auto copy = - module->GetComputationWithName("entry")->GetInstructionWithName("copy"); - EXPECT_EQ(copy->shape().layout().memory_space(), kAlternateMemorySpace); - auto neg1_operand = module->GetComputationWithName("true_computation2") - ->GetInstructionWithName("neg1") - ->operand(0); - auto neg2_operand = module->GetComputationWithName("false_computation2") - ->GetInstructionWithName("neg2") - ->operand(0); - auto neg3_operand = module->GetComputationWithName("false_computation1") - ->GetInstructionWithName("neg3") - ->operand(0); - EXPECT_EQ(neg1_operand->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(neg2_operand->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(neg3_operand->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, NestedConditionalBufferReuseVerificationBug) { - // Tests a spurious verification failure when there are nested conditionals - // and the innermost conditional computation reuses the buffer. Here, both the - // parameter of true_computation2 and neg2 will get the same buffer. Make sure - // that verification doesn't claim a failure in this case. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation2 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - neg1 = f32[3]{0} negate(gte) - neg2 = f32[3]{0} negate(neg1) - ROOT neg3 = f32[3]{0} negate(neg2) - } - - false_computation2 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg4 = f32[3]{0} negate(gte) - } - - true_computation1 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - slice = f32[1]{0} slice(gte), slice={[0:1]} - bitcast = f32[] bitcast(slice) - constant = f32[] constant(0.0) - compare = pred[] compare(bitcast, constant), direction=GT - tuple = (f32[3]{0}) tuple(gte) - ROOT conditional = f32[3]{0} conditional(compare, tuple, tuple), true_computation=true_computation2, false_computation=false_computation2 - } - - false_computation1 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg5 = f32[3]{0} negate(gte) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}) tuple(copy) - ROOT conditional = f32[3]{0} conditional(p1, tuple, tuple), true_computation=true_computation1, false_computation=false_computation1 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, WhileInsideNestedConditionalVerificationBug) { - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}) parameter(0) - ROOT constant = pred[] constant(true) - } - - while_body { - p0 = (f32[3]{0}) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - negate0 = f32[3]{0} negate(gte0) - ROOT tuple = (f32[3]{0}) tuple(negate0) - } - - true_computation2 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - tuple = (f32[3]{0}) tuple(gte) - while = (f32[3]{0}) while(tuple), condition=while_cond, body=while_body - while_gte0 = f32[3]{0} get-tuple-element(while), index=0 - ROOT root = f32[3]{0} negate(while_gte0) - } - - false_computation2 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg3 = f32[3]{0} negate(gte) - } - - true_computation1 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - constant = pred[] constant(true) - tuple = (f32[3]{0}) tuple(gte) - ROOT conditional = f32[3]{0} conditional(constant, tuple, tuple), true_computation=true_computation2, false_computation=false_computation2 - } - - false_computation1 { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg3 = f32[3]{0} negate(gte) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}) tuple(copy) - ROOT conditional = f32[3]{0} conditional(p1, tuple, tuple), true_computation=true_computation1, false_computation=false_computation1 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, - ConditionalComputationBufferOverlapBeforeParam) { - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT neg2 = f32[3]{0} negate(gte) - } - - false_computation { - c = f32[3]{0} constant({0.0, 1.0, 2.0}) - neg0 = f32[3]{0} negate(c) - neg1 = f32[3]{0} negate(neg0) - p0 = (f32[3]{0}) parameter(0) - gte = f32[3]{0} get-tuple-element(p0), index=0 - ROOT add = f32[3]{0} add(gte, neg1) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}) tuple(copy) - ROOT conditional = f32[3]{0} conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace(module.get()); - - auto get_offset = [&](absl::string_view hlo_name) { - for (const auto& chunk : preset_assignments->chunks()) { - if (chunk.first.instruction->name() == hlo_name) { - return chunk.second.offset; - } - } - return static_cast(-1); - }; - - int64_t copy_offset = get_offset("copy"); - int64_t neg0_offset = get_offset("neg0"); - EXPECT_NE(copy_offset, -1); - EXPECT_NE(neg0_offset, -1); - EXPECT_NE(copy_offset, neg0_offset); -} - -TEST_F(MemorySpaceAssignmentTest, - RequestIdentifierShouldNotBeAllocatedInAlternateMem) { - // Ensure that request identifier returned by Send/Recv HLOs are not allocated - // in the alternate memory. - absl::string_view hlo_string = R"hlo( - HloModule SendRecv, is_scheduled=true - - ENTRY %AddDependency (p: f32[3]) -> f32[3] { - %p = f32[3]{0} parameter(0) - %after-all = token[] after-all() - %recv.4 = (f32[3]{0}, u32[], token[]) recv(token[] %after-all), channel_id=7 - %recv-done.4 = (f32[3]{0}, token[]) recv-done((f32[3]{0}, u32[], token[]) %recv.4), channel_id=7 - %token.1 = token[] get-tuple-element((f32[3]{0}, token[]) %recv-done.4), index=1 - %data = f32[3]{0} get-tuple-element((f32[3]{0}, token[]) %recv-done.4), index=0 - %send = (f32[3]{0}, u32[], token[]) send(f32[3]{0} %data, token[] %token.1), channel_id=2 - %send-done = token[] send-done((f32[3]{0}, u32[], token[]) %send), channel_id=2 - ROOT %add = f32[3]{0} add(f32[3]{0} %p, f32[3]{0} %data) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - for (const HloInstruction* instruction : - module->entry_computation()->instructions()) { - if (instruction->opcode() == HloOpcode::kSend || - instruction->opcode() == HloOpcode::kRecv) { - const Shape& request_identifier_shape = - ShapeUtil::GetSubshape(instruction->shape(), {1}); - EXPECT_NE(request_identifier_shape.layout().memory_space(), - kAlternateMemorySpace); - } - } -} - -TEST_F(MemorySpaceAssignmentTest, SendDoneShouldHaveSendOperand) { - // Ensure that SendDone has only a Send operand. - absl::string_view hlo_string = R"hlo( - HloModule SendRecv, is_scheduled=true - - ENTRY %AddDependency (p: f32[3]) -> f32[3] { - %p0 = f32[3]{0} parameter(0) - %p1 = f32[3]{0} parameter(1) - %neg0 = f32[3]{0} negate(f32[3]{0} %p1) - %neg1 = f32[3]{0} negate(f32[3]{0} %neg0) - %neg2 = f32[3]{0} negate(f32[3]{0} %neg1) - %neg3 = f32[3]{0} negate(f32[3]{0} %neg2) - %neg4 = f32[3]{0} negate(f32[3]{0} %neg3) - %neg5 = f32[3]{0} negate(f32[3]{0} %neg4) - %neg6 = f32[3]{0} negate(f32[3]{0} %neg5) - %after-all = token[] after-all() - %send = (f32[3]{0}, u32[], token[]) send(f32[3]{0} %p0, token[] %after-all), channel_id=2 - %send-done = token[] send-done((f32[3]{0}, u32[], token[]) %send), channel_id=2 - ROOT %add = f32[3]{0} add(f32[3]{0} %p0, f32[3]{0} %neg6) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); -} - -TEST_F(MemorySpaceAssignmentTest, SendAndSendDoneShouldGetSameAllocation) { - // Ensure that Send and SendDone have the same allocation. - absl::string_view hlo_string = R"hlo( - HloModule SendRecv, is_scheduled=true - - ENTRY %AddDependency (p: f32[3]) -> f32[3] { - %p0 = f32[3]{0} parameter(0) - %p1 = f32[3]{0} parameter(1) - %after-all = token[] after-all() - %send = (f32[3]{0}, u32[], token[]) send(f32[3]{0} %p0, token[] %after-all), channel_id=2 - %neg0 = f32[3]{0} negate(f32[3]{0} %p1) - %neg1 = f32[3]{0} negate(f32[3]{0} %neg0) - %neg2 = f32[3]{0} negate(f32[3]{0} %neg1) - %neg3 = f32[3]{0} negate(f32[3]{0} %neg2) - %neg4 = f32[3]{0} negate(f32[3]{0} %neg3) - %neg5 = f32[3]{0} negate(f32[3]{0} %neg4) - %neg6 = f32[3]{0} negate(f32[3]{0} %neg5) - %send-done = token[] send-done((f32[3]{0}, u32[], token[]) %send), channel_id=2 - ROOT %add = f32[3]{0} add(f32[3]{0} %p0, f32[3]{0} %neg6) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/4); -} - -TEST_F(MemorySpaceAssignmentTest, LastUseOpt) { - // Test that checks the last use optimization. It uses two buffers that should - // be placed in alternate memory. - // - // +-------+ - // / \ - // add1--->sub1 +-------->mul2 - // mul1===>add2 - // - // Without the last use optimization, the mul1 buffer will be assigned first - // (because it is larger) to offset 0. Then, add1 will be scheduled for the - // add1 to sub1 segment. Because offset 0 is available, it will get that - // offset. But because offset 0 is not available in the sub1 to mul2 offset, - // it will end up in unnecessary copies. With the last use optimization, these - // copies can be optimized away. - HloComputation::Builder builder(TestName()); - Shape shape1 = ShapeUtil::MakeShape(F32, {2, 3}); - Shape shape2 = ShapeUtil::MakeShape(F32, {2, 4}); - PaddingConfig padding_config = MakeEdgePaddingConfig({{0, 0}, {0, 1}}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape1, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape2, "p1")); - HloInstruction* add1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape1, HloOpcode::kAdd, p0, p0)); - HloInstruction* sub1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape1, HloOpcode::kSubtract, p0, add1)); - HloInstruction* mul1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape2, HloOpcode::kMultiply, p1, p1)); - HloInstruction* add2 = builder.AddInstruction( - HloInstruction::CreateBinary(shape2, HloOpcode::kAdd, mul1, p1)); - HloInstruction* mul2 = builder.AddInstruction( - HloInstruction::CreateBinary(shape1, HloOpcode::kMultiply, add1, sub1)); - HloInstruction* padding_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::Zero(F32))); - HloInstruction* padded_mul2 = builder.AddInstruction( - HloInstruction::CreatePad(shape2, mul2, padding_value, padding_config)); - HloInstruction* add3 = builder.AddInstruction( - HloInstruction::CreateBinary(shape2, HloOpcode::kAdd, add2, padded_mul2)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, add1, sub1, mul1, add2, mul2, - padding_value, padded_mul2, add3}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - EXPECT_THAT( - mul2, - op::Multiply( - op::Add(op::Parameter(0), op::Parameter(0)), - op::Subtract(op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::Parameter(0)), - op::Add(op::Parameter(0), op::Parameter(0))))); -} - -TEST_F(MemorySpaceAssignmentTest, NonEntryComputationSchedule1) { - // Test to ensure CopyStart/CopyDone is placed only in the entry computation. - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, scalar_shape}); - - auto cond_builder = HloComputation::Builder("WhileCond"); - // Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element) - HloInstruction* cond_param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "cond_param")); - HloInstruction* cond_iter = cond_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1)); - HloInstruction* cond_limit = cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(50.f))); - // Free cond_param[] (16 bytes), Alloc PRED[] (1 byte) - HloInstruction* cond_lt = cond_builder.AddInstruction( - HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), cond_iter, - cond_limit, ComparisonDirection::kLt)); - HloComputation* cond_computation = - module->AddEmbeddedComputation(cond_builder.Build()); - - auto body_builder = HloComputation::Builder("WhileBody"); - // Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element) - HloInstruction* body_param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "body_param")); - HloInstruction* body_iter = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, body_param, 1)); - HloInstruction* body_data = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, body_param, 0)); - HloInstruction* body_iter_increment = body_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.f))); - HloInstruction* body_iter_next = - body_builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape, HloOpcode::kAdd, body_iter, body_iter_increment)); - HloInstruction* body_data_increment = - body_builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.f, 2.f, 3.f}, {4.f, 5.f, 6.f}}))); - HloInstruction* body_data_mul = - body_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kMultiply, body_data, body_data)); - HloInstruction* body_data_add = - body_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, body_data, body_data_increment)); - HloInstruction* body_data_next = - body_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, body_data_add, body_data_mul)); - HloInstruction* body_out = body_builder.AddInstruction( - HloInstruction::CreateTuple({body_data_next, body_iter_next})); - HloComputation* body_computation = - module->AddEmbeddedComputation(body_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* data = builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "param_iter")); - HloInstruction* iter = builder.AddInstruction( - HloInstruction::CreateParameter(1, scalar_shape, "param_data")); - HloInstruction* p2 = - builder.AddInstruction(HloInstruction::CreateParameter(2, shape, "p2")); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({data, iter})); - HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( - tuple_shape, cond_computation, body_computation, tuple)); - HloInstruction* while_data = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, while_op, 0)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, while_data, p2)); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(cond_computation, - {cond_param, cond_iter, cond_limit, cond_lt}); - schedule.set_sequence(body_computation, - {body_param, body_iter, body_data, body_iter_increment, - body_iter_next, body_data_increment, body_data_mul, - body_data_add, body_data_next, body_out}); - schedule.set_sequence(entry_computation, - {iter, data, p2, tuple, while_op, while_data, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 50); -} - -TEST_F(MemorySpaceAssignmentTest, NonEntryComputationSchedule2) { - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape shape2 = ShapeUtil::MakeShape(xla::F32, {3, 3}); - - auto call_builder = HloComputation::Builder("Call"); - HloInstruction* call_param = call_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "call_param")); - HloInstruction* call_param2 = call_builder.AddInstruction( - HloInstruction::CreateParameter(1, shape2, "call_param2")); - HloInstruction* slice = call_builder.AddInstruction( - HloInstruction::CreateSlice(shape, call_param2, {0, 0}, {2, 3}, {1, 1})); - HloInstruction* mul = - call_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kMultiply, call_param, slice)); - HloInstruction* negate0 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, mul)); - HloInstruction* negate1 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* add0 = - call_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, call_param, negate7)); - HloComputation* call_computation = - module->AddEmbeddedComputation(call_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape2, "p1")); - HloInstruction* add1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, p0)); - HloInstruction* add2 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add1, p0)); - HloInstruction* negate8 = builder.AddInstruction( - HloInstruction::CreateUnary(shape2, HloOpcode::kNegate, p1)); - HloInstruction* call = builder.AddInstruction( - HloInstruction::CreateCall(shape, {add1, negate8}, call_computation)); - HloInstruction* add3 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, add1)); - HloInstruction* add4 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, call, add3)); - HloInstruction* add5 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add2, add4)); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - call_computation, - {call_param, call_param2, slice, mul, negate0, negate1, negate2, negate3, - negate4, negate5, negate6, negate7, add0}); - schedule.set_sequence(entry_computation, - {p0, p1, add1, add2, negate8, call, add3, add4, add5}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 5); -} - -TEST_F(MemorySpaceAssignmentTest, NonEntryComputationSchedule3) { - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape shape2 = ShapeUtil::MakeShape(xla::F32, {3, 3}); - - auto call_builder = HloComputation::Builder("Call"); - HloInstruction* call_param = call_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "call_param")); - // Use shape2 here which is larger (scheduled earlier) to occupy alternate - // memory at the beginning. This should cause a situation where the prefetch - // of add1 later in the function body gets the wrong offset which cannot be - // communicated to the outside the function. - HloInstruction* iota = - call_builder.AddInstruction(HloInstruction::CreateIota(shape2, 0)); - HloInstruction* slice = call_builder.AddInstruction( - HloInstruction::CreateSlice(shape, iota, {0, 0}, {2, 3}, {1, 1})); - HloInstruction* mul = - call_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kMultiply, call_param, slice)); - HloInstruction* negate0 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, mul)); - HloInstruction* negate1 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = call_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* add0 = - call_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, call_param, negate7)); - HloComputation* call_computation = - module->AddEmbeddedComputation(call_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* add1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, p0)); - HloInstruction* add2 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add1, p0)); - HloInstruction* call = builder.AddInstruction( - HloInstruction::CreateCall(shape, {add1}, call_computation)); - HloInstruction* add3 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, call, add1)); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - call_computation, - {call_param, iota, slice, mul, negate0, negate1, negate2, negate3, - negate4, negate5, negate6, negate7, add0}); - schedule.set_sequence(entry_computation, {p0, add1, add2, call, add3}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 5); -} - -// TODO(berkin): This might be an incorrect input graph, investigate. -TEST_F(MemorySpaceAssignmentTest, DISABLED_NonEntryComputationSchedule4) { - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape shape2 = ShapeUtil::MakeShape(xla::F32, {3, 3}); - - auto true_builder = HloComputation::Builder("True"); - HloInstruction* true_param = true_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "true_param")); - HloInstruction* iota = - true_builder.AddInstruction(HloInstruction::CreateIota(shape2, 0)); - HloInstruction* slice = true_builder.AddInstruction( - HloInstruction::CreateSlice(shape, iota, {0, 0}, {2, 3}, {1, 1})); - HloInstruction* mul = - true_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kMultiply, true_param, slice)); - HloInstruction* negate0 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, mul)); - HloInstruction* negate1 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = true_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* add0 = - true_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, true_param, negate7)); - HloComputation* true_computation = - module->AddEmbeddedComputation(true_builder.Build()); - - auto false_builder = HloComputation::Builder("False"); - HloInstruction* false_param = false_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "false_param")); - HloComputation* false_computation = - module->AddEmbeddedComputation(false_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* add1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, p0)); - HloInstruction* add2 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add1, p0)); - HloInstruction* pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); - HloInstruction* conditional = - builder.AddInstruction(HloInstruction::CreateConditional( - shape, pred, add1, true_computation, add2, false_computation)); - HloInstruction* add3 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, conditional, add1)); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - true_computation, - {true_param, iota, slice, mul, negate0, negate1, negate2, negate3, - negate4, negate5, negate6, negate7, add0}); - schedule.set_sequence(false_computation, {false_param}); - schedule.set_sequence(entry_computation, - {p0, add1, add2, pred, conditional, add3}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 5); -} - -TEST_F(MemorySpaceAssignmentTest, NonEntryComputationSchedule5) { - // This test reproduces the failure in b/143288178. Given a graph like the - // following: - // - // ... = foo(a) - // tuple = tuple((..., a) - // ... = while(tuple) { - // p = param(0) - // a1 = get-tuple-element(p), index=n-1 - // ... - // ROOT tuple((..., a1)) - // } - // - // If a copy to alternate memory is inserted before foo, and if the size of - // the while body is less than max prefetch interval so that the copy-done is - // kept in the alternate memory, then we end up referring to the copy-done in - // the root instruction of the while loop body. I.e., - // - // cs = copy-start(a) - // ... - // cd = copy-done(cs) - // ... = foo(cd) - // tuple = tuple((..., cd) - // ... = while(tuple) { - // p = param(0) - // a1 = get-tuple-element(p), index=n-1 - // ... - // ROOT tuple((..., cd)) <-- Error: cd belongs to outside computation. - // } - // - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); - Shape tuple_shape = - ShapeUtil::MakeTupleShape({shape, scalar_shape, scalar_shape}); - - auto cond_builder = HloComputation::Builder("WhileCond"); - HloInstruction* cond_param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "cond_param")); - HloInstruction* cond_iter = cond_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1)); - HloInstruction* cond_limit = cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(50.f))); - HloInstruction* cond_lt = cond_builder.AddInstruction( - HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), cond_iter, - cond_limit, ComparisonDirection::kLt)); - HloComputation* cond_computation = - module->AddEmbeddedComputation(cond_builder.Build()); - - auto body_builder = HloComputation::Builder("WhileBody"); - HloInstruction* body_param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "body_param")); - HloInstruction* body_iter = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, body_param, 1)); - HloInstruction* body_data = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, body_param, 0)); - HloInstruction* body_iter_increment = body_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.f))); - HloInstruction* body_iter_next = - body_builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape, HloOpcode::kAdd, body_iter, body_iter_increment)); - HloInstruction* body_data2 = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, body_param, 2)); - HloInstruction* body_out = body_builder.AddInstruction( - HloInstruction::CreateTuple({body_data, body_iter_next, body_data2})); - HloComputation* body_computation = - module->AddEmbeddedComputation(body_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* data = builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "param_data")); - HloInstruction* iter = builder.AddInstruction( - HloInstruction::CreateParameter(1, scalar_shape, "param_iter")); - HloInstruction* data2 = builder.AddInstruction( - HloInstruction::CreateParameter(2, scalar_shape, "param_data2")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, data)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* sub = builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape, HloOpcode::kSubtract, iter, data2)); - HloInstruction* tuple = builder.AddInstruction( - HloInstruction::CreateTuple({negate7, iter, data2})); - HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( - tuple_shape, cond_computation, body_computation, tuple)); - HloInstruction* while_data = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, while_op, 1)); - HloInstruction* root = - builder.AddInstruction(HloInstruction::CreateTuple({while_data, sub})); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(cond_computation, - {cond_param, cond_iter, cond_limit, cond_lt}); - schedule.set_sequence(body_computation, - {body_param, body_iter, body_data, body_iter_increment, - body_iter_next, body_data2, body_out}); - schedule.set_sequence( - entry_computation, - {iter, data, data2, negate0, negate1, negate2, negate3, negate4, negate5, - negate6, negate7, sub, tuple, while_op, while_data, root}); - CHECK_OK(module->set_schedule(schedule)); - - // Set a large max prefetch interval so that the buffer can be kept in - // alternate memory. - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 20); -} - -TEST_F(MemorySpaceAssignmentTest, NonEntryComputationSchedule6) { - auto module = CreateNewVerifiedModule(); - Shape shape = ShapeUtil::MakeShape(xla::F32, {2, 3}); - Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, scalar_shape, shape}); - - auto cond_builder = HloComputation::Builder("WhileCond"); - HloInstruction* cond_param = cond_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "cond_param")); - HloInstruction* cond_iter = cond_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1)); - HloInstruction* cond_limit = cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(50.f))); - HloInstruction* cond_lt = cond_builder.AddInstruction( - HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), cond_iter, - cond_limit, ComparisonDirection::kLt)); - HloComputation* cond_computation = - module->AddEmbeddedComputation(cond_builder.Build()); - - auto body_builder = HloComputation::Builder("WhileBody"); - HloInstruction* body_param = body_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "body_param")); - HloInstruction* body_iter = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(scalar_shape, body_param, 1)); - HloInstruction* body_data = body_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, body_param, 0)); - HloInstruction* body_negate0 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_data)); - HloInstruction* body_negate1 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate0)); - HloInstruction* body_negate2 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate1)); - HloInstruction* body_negate3 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate2)); - HloInstruction* body_negate4 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate3)); - HloInstruction* body_negate5 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate4)); - HloInstruction* body_negate6 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate5)); - HloInstruction* body_negate7 = body_builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, body_negate6)); - HloInstruction* body_iter_increment = body_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.f))); - HloInstruction* body_iter_next = - body_builder.AddInstruction(HloInstruction::CreateBinary( - scalar_shape, HloOpcode::kAdd, body_iter, body_iter_increment)); - HloInstruction* body_out = body_builder.AddInstruction( - HloInstruction::CreateTuple({body_data, body_iter_next, body_negate7})); - HloComputation* body_computation = - module->AddEmbeddedComputation(body_builder.Build()); - - auto builder = HloComputation::Builder(TestName()); - HloInstruction* data = builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "param_data")); - HloInstruction* iter = builder.AddInstruction( - HloInstruction::CreateParameter(1, scalar_shape, "param_iter")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, data)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* negate7 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate6)); - HloInstruction* tuple = builder.AddInstruction( - HloInstruction::CreateTuple({data, iter, negate7})); - HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( - tuple_shape, cond_computation, body_computation, tuple)); - HloInstruction* while_data = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, while_op, 0)); - HloInstruction* while_data2 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, while_op, 2)); - HloInstruction* root = builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, while_data, while_data2)); - HloComputation* entry_computation = - module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(cond_computation, - {cond_param, cond_iter, cond_limit, cond_lt}); - schedule.set_sequence( - body_computation, - {body_param, body_iter, body_data, body_negate0, body_negate1, - body_negate2, body_negate3, body_negate4, body_negate5, body_negate6, - body_negate7, body_iter_increment, body_iter_next, body_out}); - schedule.set_sequence( - entry_computation, - {iter, data, negate0, negate1, negate2, negate3, negate4, negate5, - negate6, negate7, tuple, while_op, while_data, while_data2, root}); - CHECK_OK(module->set_schedule(schedule)); - - // Pick a large max prefetch interval to ensure all the while inputs are - // allocated in the alternate memory. - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/25); - - // Index {0} of the while loop argument is not written inside the while loop, - // so it can be trivially placed in the alternate memory space. - *ShapeUtil::GetMutableSubshape(&tuple_shape, {0})->mutable_layout() = - LayoutUtil::MakeLayout( - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID, - /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, - /*element_size_in_bits=*/0, kAlternateMemorySpace); - // Index {1} is a scalar, so it is always placed in the default memory. - *ShapeUtil::GetMutableSubshape(&tuple_shape, {1})->mutable_layout() = - LayoutUtil::MakeLayout( - /*minor_to_major=*/{}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID, - /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, - /*element_size_in_bits=*/0, kDefaultMemorySpace); - // Index {2} of the while loop is placed in the default memory. - *ShapeUtil::GetMutableSubshape(&tuple_shape, {2})->mutable_layout() = - LayoutUtil::MakeLayout( - /*minor_to_major=*/{1, 0}, - /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, - /*index_primitive_type=*/PRIMITIVE_TYPE_INVALID, - /*pointer_primitive_type=*/PRIMITIVE_TYPE_INVALID, - /*element_size_in_bits=*/0, kDefaultMemorySpace); - - // Expect the layout for the while loop and its aliased buffers. - EXPECT_THAT(while_op, op::ShapeWithLayout(tuple_shape)); - EXPECT_THAT(while_op->operand(0), op::ShapeWithLayout(tuple_shape)); - EXPECT_THAT(cond_param, op::ShapeWithLayout(tuple_shape)); - EXPECT_THAT(body_param, op::ShapeWithLayout(tuple_shape)); - EXPECT_THAT(body_out, op::ShapeWithLayout(tuple_shape)); -} - -TEST_F(MemorySpaceAssignmentTest, DanglingCopy) { - // This situation was encountered in vss, where there is a mismatch in the - // memory space in preset assignments and the output graph. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - - HloInstruction* p = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* p0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 0)); - HloInstruction* p1a = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 1)); - HloInstruction* copy = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kCopy, p1a)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* p1b = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 1)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1b)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, {p, p0, negate0, negate1, negate2, negate3, negate4, negate5, - negate6, p1a, copy, p1b, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, MultiOutputFusion) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - auto module = CreateNewVerifiedModule(); - - HloComputation::Builder fusion_builder("fusion"); - HloInstruction* fusion_param0 = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion_param1 = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(1, shape, "p1")); - fusion_builder.AddInstruction( - HloInstruction::CreateTuple({fusion_param0, fusion_param1})); - HloComputation* fusion_computation = - module->AddEmbeddedComputation(fusion_builder.Build()); - - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion = builder.AddInstruction(HloInstruction::CreateFusion( - tuple_shape, HloInstruction::FusionKind::kCustom, {p0, p0}, - fusion_computation)); - HloInstruction* element0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion, 0)); - HloInstruction* element1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion, 1)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, element0, element1)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, fusion, element0, element1, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, TupleInput) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - auto module = CreateNewVerifiedModule(); - - HloComputation::Builder fusion_builder("fusion"); - HloInstruction* fusion_param = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* fusion_element0 = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion_param, 0)); - HloInstruction* fusion_element1 = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion_param, 1)); - fusion_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, fusion_element0, fusion_element1)); - HloComputation* fusion_computation = - module->AddEmbeddedComputation(fusion_builder.Build()); - - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p1)); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({negate0, negate1})); - HloInstruction* fusion = builder.AddInstruction(HloInstruction::CreateFusion( - shape, HloInstruction::FusionKind::kCustom, {tuple}, fusion_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, tuple, fusion}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, TupleToTuple1) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - auto module = CreateNewVerifiedModule(); - - HloComputation::Builder fusion0_builder("fusion0"); - HloInstruction* fusion0_param0 = fusion0_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion0_param1 = fusion0_builder.AddInstruction( - HloInstruction::CreateParameter(1, shape, "p1")); - fusion0_builder.AddInstruction( - HloInstruction::CreateTuple({fusion0_param0, fusion0_param1})); - HloComputation* fusion0_computation = - module->AddEmbeddedComputation(fusion0_builder.Build()); - - HloComputation::Builder fusion1_builder("fusion1"); - HloInstruction* fusion1_param = fusion1_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* fusion1_element0 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion1_param, 0)); - HloInstruction* fusion1_element1 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion1_param, 1)); - fusion1_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, fusion1_element0, fusion1_element1)); - HloComputation* fusion1_computation = - module->AddEmbeddedComputation(fusion1_builder.Build()); - - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion0 = builder.AddInstruction(HloInstruction::CreateFusion( - tuple_shape, HloInstruction::FusionKind::kCustom, {p0, p0}, - fusion0_computation)); - HloInstruction* element0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion0, 0)); - HloInstruction* element1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion0, 1)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add0 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, element0, element1)); - HloInstruction* add1 = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, add0, negate6)); - HloInstruction* fusion1 = builder.AddInstruction( - HloInstruction::CreateFusion(shape, HloInstruction::FusionKind::kCustom, - {fusion0}, fusion1_computation)); - HloInstruction* mul = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, add1, fusion1)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, - {p0, fusion0, element0, element1, negate0, negate1, negate2, negate3, - negate4, negate5, negate6, add0, add1, fusion1, mul}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 5); - EXPECT_THAT(fusion1, - op::Fusion(op::Tuple( - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::Fusion(), 0)), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::Fusion(), 1))))); -} - -TEST_F(MemorySpaceAssignmentTest, TupleToTuple2) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - Shape nested_tuple_shape = ShapeUtil::MakeTupleShape({shape, tuple_shape}); - auto module = CreateNewVerifiedModule(); - - HloComputation::Builder fusion0_builder("fusion0"); - HloInstruction* fusion0_param0 = fusion0_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion0_param1 = fusion0_builder.AddInstruction( - HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* fusion0_tuple = fusion0_builder.AddInstruction( - HloInstruction::CreateTuple({fusion0_param0, fusion0_param1})); - fusion0_builder.AddInstruction( - HloInstruction::CreateTuple({fusion0_param0, fusion0_tuple})); - HloComputation* fusion0_computation = - module->AddEmbeddedComputation(fusion0_builder.Build()); - - HloComputation::Builder fusion1_builder("fusion1"); - HloInstruction* fusion1_param = fusion1_builder.AddInstruction( - HloInstruction::CreateParameter(0, nested_tuple_shape, "p")); - HloInstruction* fusion1_element0 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion1_param, 0)); - HloInstruction* fusion1_element1 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(tuple_shape, fusion1_param, 1)); - HloInstruction* fusion1_element2 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion1_element1, 1)); - fusion1_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, fusion1_element0, fusion1_element2)); - HloComputation* fusion1_computation = - module->AddEmbeddedComputation(fusion1_builder.Build()); - - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion0 = builder.AddInstruction(HloInstruction::CreateFusion( - nested_tuple_shape, HloInstruction::FusionKind::kCustom, {p0, p0}, - fusion0_computation)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* fusion1 = builder.AddInstruction( - HloInstruction::CreateFusion(shape, HloInstruction::FusionKind::kCustom, - {fusion0}, fusion1_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, {p0, fusion0, negate0, negate1, negate2, negate3, negate4, - negate5, negate6, fusion1}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), 5); - - EXPECT_THAT( - fusion1, - op::Fusion(op::Tuple( - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::Fusion(), 0)), - op::Tuple( - op::AsyncCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::GetTupleElement(op::Fusion(), 1), 0)), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement( - op::GetTupleElement(op::Fusion(), 1), 1)))))); -} - -TEST_F(MemorySpaceAssignmentTest, TupleToTuple3) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - auto module = CreateNewVerifiedModule(); - - HloComputation::Builder fusion0_builder("fusion0"); - HloInstruction* fusion0_param0 = fusion0_builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion0_param1 = fusion0_builder.AddInstruction( - HloInstruction::CreateParameter(1, shape, "p1")); - fusion0_builder.AddInstruction( - HloInstruction::CreateTuple({fusion0_param0, fusion0_param1})); - HloComputation* fusion0_computation = - module->AddEmbeddedComputation(fusion0_builder.Build()); - - HloComputation::Builder fusion1_builder("fusion1"); - HloInstruction* fusion1_param = fusion1_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* fusion1_element0 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion1_param, 0)); - HloInstruction* fusion1_element1 = fusion1_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, fusion1_param, 1)); - fusion1_builder.AddInstruction(HloInstruction::CreateBinary( - shape, HloOpcode::kAdd, fusion1_element0, fusion1_element1)); - HloComputation* fusion1_computation = - module->AddEmbeddedComputation(fusion1_builder.Build()); - - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* fusion0 = builder.AddInstruction(HloInstruction::CreateFusion( - tuple_shape, HloInstruction::FusionKind::kCustom, {p0, p0}, - fusion0_computation)); - HloInstruction* fusion1 = builder.AddInstruction( - HloInstruction::CreateFusion(shape, HloInstruction::FusionKind::kCustom, - {fusion0}, fusion1_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, fusion0, fusion1}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - EXPECT_THAT(fusion1, op::Fusion(op::Fusion())); -} - -TEST_F(MemorySpaceAssignmentTest, InputOutputAlias) { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape tuple_shape = ShapeUtil::MakeTupleShape({shape, shape}); - HloInstruction* p = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p")); - HloInstruction* p0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 0)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(shape, p, 1)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - HloInstruction* negate7 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, add)); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({p0, add})); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, {p, p0, negate0, negate1, negate2, negate3, negate4, negate5, - negate6, p1, add, negate7, tuple}); - CHECK_OK(module->set_schedule(schedule)); - - // Make input {0} alias with output {0} and input {1} alias with output {1}. - CHECK_OK(module->input_output_alias_config().SetUpAlias({0}, 0, {0})); - CHECK_OK(module->input_output_alias_config().SetUpAlias({1}, 0, {1})); - - AssignMemorySpace(module.get()); - - // Make sure the input is in the default memory space. - EXPECT_EQ(p->shape().tuple_shapes(0).layout().memory_space(), - kDefaultMemorySpace); - EXPECT_EQ(p->shape().tuple_shapes(1).layout().memory_space(), - kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, CostAnalysis) { - // This is mostly a smoke test since it's difficult and brittle to work out - // the cost of the HLO instructions. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, - negate3, negate4, negate5, negate6, add}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpaceUsingCostAnalysis(module.get()); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - // Negate instructions are in the alternate memory space (1). - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - EXPECT_THAT(negate0, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate1, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate2, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate3, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate4, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate5, op::ShapeWithLayout(shape_in_alternate_mem)); - EXPECT_THAT(negate6, op::ShapeWithLayout(shape_in_alternate_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, MemoryBoundednessBufferIntervalCompare) { - // This test is carefully crafted to force only negates to be allocated to the - // alternate memory. The graph consists of interleaving negate and tanh - // operations: - // - // +------+ +-------+ +----- - // / \ / \ / - // negate tanh negate tanh negate tanh - // \ / \ / - // +--------+ +---------+ - // - // The alternate memory is sized to fit only two f32[4,3] tensors at a time. - // Also, transcendentals are made to be lower bandwidth than FLOPs. So, the - // MemoryBoundednessBufferIntervalCompare should prioritize the negates, which - // are more memory bound. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {4, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* tanh0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p1)); - HloInstruction* tanh1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, tanh0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* tanh2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, tanh1)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* tanh3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, tanh2)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* tanh4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, tanh3)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({tanh4, negate4})); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, - {p0, p1, tanh0, negate0, tanh1, negate1, tanh2, negate2, - tanh3, negate3, tanh4, negate4, tuple}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpaceUsingCostAnalysis(module.get()); - // Parameters are in the default memory space. - EXPECT_THAT(p0, op::ShapeWithLayout(shape)); - EXPECT_THAT(p1, op::ShapeWithLayout(shape)); - Shape shape_in_default_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {4, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kDefaultMemorySpace); - // Expect only negates to be in alternate memory space. Not all might fit but - // make sure at least one does. - std::vector negate_instructions = {negate0, negate1, negate2, - negate3, negate4}; - int64_t num_negates_in_alternate_mem = absl::c_count_if( - negate_instructions, [&](const HloInstruction* instruction) { - return instruction->shape().layout().memory_space() == - kAlternateMemorySpace; - }); - EXPECT_GE(num_negates_in_alternate_mem, 1); - EXPECT_THAT(tanh0, op::ShapeWithLayout(shape_in_default_mem)); - EXPECT_THAT(tanh1, op::ShapeWithLayout(shape_in_default_mem)); - EXPECT_THAT(tanh2, op::ShapeWithLayout(shape_in_default_mem)); - EXPECT_THAT(tanh3, op::ShapeWithLayout(shape_in_default_mem)); - EXPECT_THAT(tanh4, op::ShapeWithLayout(shape_in_default_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, TestColorOutputInAlternateMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(negate0) - negate1 = f32[3,4]{1,0} add(negate0, p1) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} add(negate1, p1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} add(negate2, p1) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} add(negate3, tanh0) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* tanh0 = FindInstruction(module.get(), "tanh0"); - HloPosition tanh0_position{tanh0, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {tanh0_position, kAlternateMemorySpace}}; - - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "tanh0" } - override_options { assign_last: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), std::nullopt, std::nullopt, - msa_sort_order_overrides); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - EXPECT_EQ( - FindInstruction(module.get(), "tanh0")->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ColorInputInAlternateMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(negate0) - negate1 = f32[3,4]{1,0} add(negate0, p1) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} add(negate1, p1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} add(negate2, p1) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} add(negate3, tanh0) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* negate4 = FindInstruction(module.get(), "negate4"); - HloUse tanh0_use_during_negate4{negate4, 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {tanh0_use_during_negate4, kAlternateMemorySpace}}; - - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "tanh0" } - override_options { assign_last: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), std::nullopt, std::nullopt, - msa_sort_order_overrides); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - EXPECT_EQ(FindInstruction(module.get(), "negate4") - ->operand(1) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - ColorOutputInDefaultMemoryUseInAlternateMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(negate0) - negate1 = f32[3,4]{1,0} add(negate0, p1) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} add(negate1, p1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} add(negate2, p1) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} add(negate3, tanh0) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* tanh0 = FindInstruction(module.get(), "tanh0"); - HloPosition tanh0_position{tanh0, {}}; - HloInstruction* negate4 = FindInstruction(module.get(), "negate4"); - HloUse tanh0_use_during_negate4{negate4, 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {tanh0_position, kDefaultMemorySpace}, - {tanh0_use_during_negate4, kAlternateMemorySpace}}; - - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "tanh0" } - override_options { assign_last: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), std::nullopt, std::nullopt, - msa_sort_order_overrides); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - EXPECT_EQ( - FindInstruction(module.get(), "tanh0")->shape().layout().memory_space(), - kDefaultMemorySpace); - EXPECT_EQ(FindInstruction(module.get(), "negate4") - ->operand(1) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - ColorOutputInAlternateMemoryUseInDefaultMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { -p0 = f32[3,4]{1,0} parameter(0) -p1 = f32[3,4]{1,0} parameter(1) -tanh0 = f32[3,4]{1,0} tanh(p0) -negate0 = f32[3,4]{1,0} negate(p1) -tanh1 = f32[3,4]{1,0} tanh(negate0) -negate1 = f32[3,4]{1,0} add(negate0, p1) -tanh2 = f32[3,4]{1,0} tanh(tanh1) -negate2 = f32[3,4]{1,0} add(negate1, p1) -tanh3 = f32[3,4]{1,0} tanh(tanh2) -negate3 = f32[3,4]{1,0} add(negate2, p1) -tanh4 = f32[3,4]{1,0} tanh(tanh3) -negate4 = f32[3,4]{1,0} add(negate3, tanh0) -ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* tanh0 = FindInstruction(module.get(), "tanh0"); - HloPosition tanh0_position{tanh0, {}}; - HloInstruction* negate4 = FindInstruction(module.get(), "negate4"); - HloUse tanh0_use_during_negate4{negate4, 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {tanh0_position, kAlternateMemorySpace}, - {tanh0_use_during_negate4, kDefaultMemorySpace}}; - - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "tanh0" } - override_options { assign_last: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), std::nullopt, std::nullopt, - msa_sort_order_overrides); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - EXPECT_EQ( - FindInstruction(module.get(), "tanh0")->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(FindInstruction(module.get(), "negate4") - ->operand(1) - ->shape() - .layout() - .memory_space(), - kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ColorAlternateUsesInAlternateMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { -p0 = f32[3,4]{1,0} parameter(0) -p1 = f32[3,4]{1,0} parameter(1) -tanh0 = f32[3,4]{1,0} tanh(p0) -negate0 = f32[3,4]{1,0} negate(p1) -tanh1 = f32[3,4]{1,0} tanh(negate0) -negate1 = f32[3,4]{1,0} add(negate0, p1) -tanh2 = f32[3,4]{1,0} tanh(tanh1) -negate2 = f32[3,4]{1,0} add(negate1, p1) -tanh3 = f32[3,4]{1,0} tanh(tanh2) -negate3 = f32[3,4]{1,0} add(negate2, p1) -tanh4 = f32[3,4]{1,0} tanh(tanh3) -negate4 = f32[3,4]{1,0} add(negate3, tanh0) -ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* negate0 = FindInstruction(module.get(), "negate0"); - HloPosition negate0_position{negate0, {}}; - HloInstruction* negate1 = FindInstruction(module.get(), "negate1"); - HloUse negate0_use_during_negate1{negate1, 0, {}}; - HloInstruction* tanh1 = FindInstruction(module.get(), "tanh1"); - HloUse negate0_use_during_tanh1{tanh1, 0, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate0_position, kAlternateMemorySpace}, - {negate0_use_during_tanh1, kDefaultMemorySpace}, - {negate0_use_during_negate1, kAlternateMemorySpace}}; - - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "tanh0" } - override_options { assign_first: true } - } - overrides { - hlo_position_matcher { instruction_name_regex: "negate0" } - override_options { assign_last: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), std::nullopt, std::nullopt, - msa_sort_order_overrides); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - EXPECT_EQ( - FindInstruction(module.get(), "negate0")->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(FindInstruction(module.get(), "negate1") - ->operand(0) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(FindInstruction(module.get(), "tanh1") - ->operand(0) - ->shape() - .layout() - .memory_space(), - kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ColorBuffersWithoutRepacking) { - // Without repacking, the following is the expected layout: - // ...... chunk: 3072-4096 - // #####. chunk: 2048-3072 - // #####. chunk: 1024-2048 - // .##... chunk: 0-1024 - // Since we block a chunk for coloring of negate1, we are not able to allocate - // negate3 in alternate memory without repacking. - absl::string_view hlo_string = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - /* parameters */ - p0 = f32[32,16] parameter(0) - p1 = f32[16,16] parameter(1) - p2 = f32[32,16] parameter(2) - - negate0 = f32[32,16] negate(p0) // We will set highest priority for this. - negate1 = f32[16,16] negate(p1) // We will color this, but set lowest priority. - negate2 = f32[16,16] negate(negate1) - negate3 = f32[32,16] negate(p2) // We will set second highest priority for this. - negate4 = f32[32,16] negate(negate0) - negate5 = f32[32,16] negate(negate3) - - ROOT tuple = (f32[16,16], f32[32,16], f32[32,16]) tuple(negate2, negate4, negate5) -})hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that for p1, p2, p3, and - // p4. - int priority = 100; - if (x.buffer->instruction()->name() == "negate0") { - priority = 1; - } else if (x.buffer->instruction()->name() == "negate3") { - priority = 2; - } else if (x.buffer->instruction()->name() == "negate1") { - priority = 1000; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 4 * 1024; - options.max_repacks = 0; - HloInstruction* negate1 = FindInstruction(module.get(), "negate1"); - HloPosition negate1_position{negate1, {}}; - options.allocate_colored_buffers_early = false; - options.buffer_colorings = {{negate1_position, kAlternateMemorySpace}}; - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), std::move(options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - EXPECT_EQ( - FindInstruction(module.get(), "negate1")->shape().layout().memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ( - FindInstruction(module.get(), "negate0")->shape().layout().memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "negate4") - ->operand(0) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - - // negate3 output and use are both in default memory. - EXPECT_EQ( - FindInstruction(module.get(), "negate3")->shape().layout().memory_space(), - kDefaultMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "negate5") - ->operand(0) - ->shape() - .layout() - .memory_space(), - kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ColorBuffersWithRepacking) { - // Without repacking, the following is the expected layout: - // ...... chunk: 3072-4096 - // #####. chunk: 2048-3072 - // #####. chunk: 1024-2048 - // .##... chunk: 0-1024 - - // With repacking, we move negate1 to offset 1024, and negate0 to offset 0. - // ...### chunk: 3072-4096 - // .##### chunk: 2048-3072 - // #####. chunk: 1024-2048 - // #####. chunk: 0-1024 - // Repacking allows us to allocate negate3 in alternate memory. - absl::string_view hlo_string = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - /* parameters */ - p0 = f32[32,16] parameter(0) - p1 = f32[16,16] parameter(1) - p2 = f32[32,16] parameter(2) - - negate0 = f32[32,16] negate(p0) // We will set highest priority for this. - negate1 = f32[16,16] negate(p1) // We will color this, but set lowest priority. - negate2 = f32[16,16] negate(negate1) - negate3 = f32[32,16] negate(p2) // We will set second highest priority for this. - negate4 = f32[32,16] negate(negate0) - negate5 = f32[32,16] negate(negate3) - - ROOT tuple = (f32[16,16], f32[32,16], f32[32,16]) tuple(negate2, negate4, negate5) -})hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that for p1, p2, p3, and - // p4. - int priority = 100; - if (x.buffer->instruction()->name() == "negate0") { - priority = 1; - } else if (x.buffer->instruction()->name() == "negate3") { - priority = 2; - } else if (x.buffer->instruction()->name() == "negate1") { - priority = 1000; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - - absl::flat_hash_map, int64_t> repack_map; - // Move "negate1" from offset 0 to 1024. - repack_map[{4, 0}] = 1024; - // Move "negate0" from offset 1024 to 0. - repack_map[{3, 1024}] = 0; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 4 * 1024; - options.max_repacks = 1; - options.repacker = &repacker; - HloInstruction* negate1 = FindInstruction(module.get(), "negate1"); - HloPosition negate1_position{negate1, {}}; - options.allocate_colored_buffers_early = false; - options.buffer_colorings = {{negate1_position, kAlternateMemorySpace}}; - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), std::move(options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - EXPECT_EQ( - FindInstruction(module.get(), "negate1")->shape().layout().memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ( - FindInstruction(module.get(), "negate0")->shape().layout().memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "negate4") - ->operand(0) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - - // negate3 output and use are both in alternate memory. - EXPECT_EQ( - FindInstruction(module.get(), "negate3")->shape().layout().memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "negate5") - ->operand(0) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); -} - -#define EXPECT_INSTRUCTION_IN(instruction_name, expected_memory_space) \ - { \ - HloInstruction* instruction = \ - FindInstruction(module.get(), instruction_name); \ - EXPECT_EQ(instruction->shape().layout().memory_space(), \ - expected_memory_space) \ - << "Instruction name:" << instruction_name; \ - } - -TEST_F(MemorySpaceAssignmentTest, - MemoryBoundednessOverrideSortOrderAssignFirst) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(tanh0) - negate1 = f32[3,4]{1,0} negate(negate0) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} negate(negate1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} negate(negate2) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} negate(negate3) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - // Override MSA sort order and try to assign all negates to alternate memory - // first. Alternate memory size is enough to fit 2 f32[4,3] tensors at a time. - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "negate(.*)" } - override_options { assign_first: true } - })pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/std::nullopt, - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - // Parameters are in the default memory space. - EXPECT_INSTRUCTION_IN("p0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("p1", kDefaultMemorySpace); - // Check that all negates are in alternate memory space except negate4. - // negate4 is a program output, so it has to land in default memory. - EXPECT_INSTRUCTION_IN("negate0", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate1", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate2", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate3", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate4", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh2", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh3", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh4", kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - MemoryBoundednessOverrideSortOrderAssignLast) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(tanh0) - negate1 = f32[3,4]{1,0} negate(negate0) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} negate(negate1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} negate(negate2) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} negate(negate3) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - // Override MSA sort order and try to assign all tanhs to alternate memory - // last. Alternate memory size is enough to fit 2 f32[4,3] tensors at a time. - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "tanh(.*)" } - override_options { assign_last: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/std::nullopt, - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - // Parameters are in the default memory space. - EXPECT_INSTRUCTION_IN("p0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("p1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate0", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate1", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate2", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate3", kAlternateMemorySpace); - // negate4 is a program output, so it has to land in default memory. - EXPECT_INSTRUCTION_IN("negate4", kDefaultMemorySpace); - // Check that all tanhs are in default memory space. - EXPECT_INSTRUCTION_IN("tanh0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh2", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh3", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh4", kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - MemoryBoundednessOverrideSortOrderBySizeLteAssignFirst) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[5,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[5,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(tanh0) - negate1 = f32[5,4]{1,0} negate(negate0) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[5,4]{1,0} negate(negate1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[5,4]{1,0} negate(negate2) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[5,4]{1,0} negate(negate3) - ROOT tuple = (f32[3,4]{1,0}, f32[5,4]{1,0}) tuple(tanh4, negate4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - // Override MSA sort order and try to assign all buffers with size lesser - // than or equal to 48 bytes to alternate memory first. - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { size_lte: 48 } - override_options { assign_first: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - Options memory_space_options = DefaultMemorySpaceOptions(); - // Set max size to 120 bytes, such that 2 f32[4,3] tensors can fit in - // alternate memory at the same time but not 1 f32[4,3] tensor and 1 - // f32[4,5] tensor. If the max size was 128 bytes, negate3 would be assigned - // to alternate memory. - memory_space_options.max_size_in_bytes = 120; - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - // Parameters are in the default memory space. - EXPECT_INSTRUCTION_IN("p0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("p1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate2", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate3", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate4", kDefaultMemorySpace); - // Check that all tanhs are in alternate memory space except tanh4. tanh4 - // is a program output, so it has to land in default memory. - EXPECT_INSTRUCTION_IN("tanh0", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh1", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh2", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh3", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh4", kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - MemoryBoundednessOverrideSortOrderBySizeGteAssignFirst) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[5,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[5,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(tanh0) - negate1 = f32[5,4]{1,0} negate(negate0) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[5,4]{1,0} negate(negate1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[5,4]{1,0} negate(negate2) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[5,4]{1,0} negate(negate3) - ROOT tuple = (f32[3,4]{1,0}, f32[5,4]{1,0}) tuple(tanh4, negate4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - // Override MSA sort order and try to assign all buffers with size greater - // than or equal to 80 bytes to alternate memory first. - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { size_gte: 80 } - override_options { assign_first: true } - } - )pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - Options memory_space_options = DefaultMemorySpaceOptions(); - // Set max size to 160 bytes to allow 2 f32[4,5] tensors to fit in alternate - // memory at the same time. tanh3 would not be prefetched because negate2 and - // negate3 would be in alternate memory at the same time leaving no space for - // tanh3. - memory_space_options.max_size_in_bytes = 160; - AssignMemorySpaceUsingCostAnalysis( - module.get(), std::move(memory_space_options), - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - // Parameters are in the default memory space. - EXPECT_INSTRUCTION_IN("p0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("p1", kDefaultMemorySpace); - // Check that all negates are in alternate memory space except negate4. - // negate4 is a program output, so it has to land in default memory. - EXPECT_INSTRUCTION_IN("negate0", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate1", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate2", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate3", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate4", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh2", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh3", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh4", kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - MemoryBoundednessOverrideSortOrderByUseAssignFirst) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(tanh0) - negate1 = f32[3,4]{1,0} negate(negate0) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} negate(negate1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} negate(negate2) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} negate(negate3) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - // Override MSA sort order and try to assign all negates to alternate memory - // first. Alternate memory size is enough to fit 2 f32[4,3] tensors at a time. - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { - hlo_use_filter { instruction_name_regex: "negate(.*)" } - } - override_options { assign_first: true } - })pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/std::nullopt, - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - // Parameters are in the default memory space. - EXPECT_INSTRUCTION_IN("p0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("p1", kDefaultMemorySpace); - // Check that all negates are in alternate memory space except negate4. - // negate4 is a program output, so it has to land in default memory. - EXPECT_INSTRUCTION_IN("negate0", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate1", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate2", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate3", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate4", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh2", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh3", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("tanh4", kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - MemoryBoundednessOverrideSortOrderByRandomFingerprintAssignFirst) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - tanh0 = f32[3,4]{1,0} tanh(p0) - negate0 = f32[3,4]{1,0} negate(p1) - tanh1 = f32[3,4]{1,0} tanh(tanh0) - negate1 = f32[3,4]{1,0} negate(negate0) - tanh2 = f32[3,4]{1,0} tanh(tanh1) - negate2 = f32[3,4]{1,0} negate(negate1) - tanh3 = f32[3,4]{1,0} tanh(tanh2) - negate3 = f32[3,4]{1,0} negate(negate2) - tanh4 = f32[3,4]{1,0} tanh(tanh3) - negate4 = f32[3,4]{1,0} negate(negate3) - ROOT tuple = (f32[3,4]{1,0}, f32[3,4]{1,0}) tuple(tanh4, negate4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - // Override MSA sort order and try to assign all negates to alternate memory - // first. Alternate memory size is enough to fit 2 f32[4,3] tensors at a time. - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { - hlo_random_filter { - seed: 123456 - selection_range_begin: 0. - selection_range_end: 0.5 - } - } - override_options { assign_first: true } - })pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - - // Show that with this random seed and bound values, we get tanh to pass the - // random filter while negate do not pass the filter. - // - // NOTE that all tanh functions are considered "identical" since the filter - // it only uses the "fingerprint" of the instruction which is based on the - // computation graph and does not include the name. Same for the negate - // functions. - const auto& hlo_position_matcher = - msa_sort_order_overrides.overrides(0).hlo_position_matcher(); - EXPECT_TRUE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "tanh0"))); - EXPECT_TRUE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "tanh1"))); - EXPECT_TRUE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "tanh2"))); - EXPECT_TRUE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "tanh3"))); - EXPECT_TRUE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "tanh4"))); - - EXPECT_FALSE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "negate0"))); - EXPECT_FALSE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "negate1"))); - EXPECT_FALSE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "negate2"))); - EXPECT_FALSE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "negate3"))); - EXPECT_FALSE(MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - hlo_position_matcher, *FindInstruction(module.get(), "negate4"))); - - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/std::nullopt, - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/std::nullopt, - /*optional_msa_sort_order_overrides=*/msa_sort_order_overrides); - // Parameters are in the default memory space. - EXPECT_INSTRUCTION_IN("p0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("p1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate0", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate1", kDefaultMemorySpace); - EXPECT_INSTRUCTION_IN("negate2", kDefaultMemorySpace); - // Since the max size is 128 bytes, there's enough space for negate3 to be - // assigned to alternate memory even though it doesn't have priority - EXPECT_INSTRUCTION_IN("negate3", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("negate4", kDefaultMemorySpace); - // Check that all tanh are in alternate memory space except tanh4. - // tanh4 is a program output, so it has to land in default memory. - EXPECT_INSTRUCTION_IN("tanh0", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh1", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh2", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh3", kAlternateMemorySpace); - EXPECT_INSTRUCTION_IN("tanh4", kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, SimpleWhileTupleTest) { - Shape s32 = ShapeUtil::MakeShape(xla::S32, {}); - Shape f32v1 = ShapeUtil::MakeShape(F32, {1}); - Shape t_s32_f32v1 = ShapeUtil::MakeTupleShape({s32, f32v1}); - auto module = CreateNewVerifiedModule("SimpleWhile"); - HloSchedule schedule(module.get()); - - // A simple compare-to-limit (x < 4) computation for a While. - // - // condition: - // const4[s32] -----------------------------------\ - // \ - // param[(s32,f32[4])] --- get-tuple-element[0] --- less-than - // - HloComputation* cond_computation; - { - auto builder = HloComputation::Builder("WhileCond"); - auto const4 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(4))); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, t_s32_f32v1, "x")); - auto index = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(const4->shape(), param, 0)); - auto compare = builder.AddInstruction( - HloInstruction::CreateCompare(ShapeUtil::MakeShape(PRED, {}), index, - const4, ComparisonDirection::kLt)); - cond_computation = module->AddEmbeddedComputation(builder.Build()); - schedule.set_sequence(cond_computation, {const4, param, index, compare}); - } - - // Builds a simple body computation for a While. - // - // body: - // constv[f32[1]] --------------------------------------\ - // \ - // /--- get-tuple-elementv[1] --- addv ---\ - // param[(s32,f32[1])] ---| tuple - // \--- get-tuple-elementc[0] --- addc ---/ - // / - // const1[s32] -----------------------------------------/ - // - HloComputation* body_computation; - { - auto builder = HloComputation::Builder("WhileBody"); - auto const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); - auto constv = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({1.1f}))); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, t_s32_f32v1, "x")); - auto indexc = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(const1->shape(), param, 0)); - auto addc = builder.AddInstruction(HloInstruction::CreateBinary( - indexc->shape(), HloOpcode::kAdd, indexc, const1)); - auto indexv = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(constv->shape(), param, 1)); - auto addv = builder.AddInstruction(HloInstruction::CreateBinary( - constv->shape(), HloOpcode::kAdd, indexv, constv)); - auto tuple = - builder.AddInstruction(HloInstruction::CreateTuple({addc, addv})); - body_computation = module->AddEmbeddedComputation(builder.Build()); - schedule.set_sequence(body_computation, {const1, constv, param, indexc, - addc, indexv, addv, tuple}); - } - - // This tests a simple while loop where the parameters are aliased with the - // output buffers. - auto builder = HloComputation::Builder("SimpleWhile"); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, t_s32_f32v1, "param")); - auto gte0 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(s32, param, 0)); - auto gte1 = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(f32v1, param, 1)); - auto tuple = - builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); - auto while0 = builder.AddInstruction(HloInstruction::CreateWhile( - t_s32_f32v1, cond_computation, body_computation, tuple)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - schedule.set_sequence(computation, {param, gte0, gte1, tuple, while0}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/50); - - // Ensure all parameters and while are placed in default memory. - Shape shape_in_default_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {4, 6}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kDefaultMemorySpace); - Shape s32_in_default_mem = ShapeUtil::MakeShapeWithDenseLayout( - xla::S32, {}, - /*minor_to_major=*/{}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kDefaultMemorySpace); - Shape f32v1_in_default_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {1}, - /*minor_to_major=*/{0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kDefaultMemorySpace); - Shape t_s32_f32v1_in_default_mem = - ShapeUtil::MakeTupleShape({s32_in_default_mem, f32v1_in_default_mem}); - EXPECT_THAT(param, op::ShapeWithLayout(t_s32_f32v1_in_default_mem)); - EXPECT_THAT(while0, op::ShapeWithLayout(t_s32_f32v1_in_default_mem)); -} - -TEST_F(MemorySpaceAssignmentTest, EvictionsShouldntBeDelayed) { - // This test reproduces an eviction scheduling bug where evictions to default - // memory can happen later than intended, causing memory corruption. This test - // is a variant of MemoryBoundednessBufferIntervalCompare but uses f32[4,3] - // tensors instead, so at most two tensors should fit in the alternate memory - // space at a given time. We have a number of redundant operations - // (tanh_redundant ops) that do not have users. The bug was due to - // SimplifyGraph removing dead instructions, and removing them from the - // schedule. However, the CopyStart/CopyDone insertion relies on the schedule - // indexes, so they could be inserted too late. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {4, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* tanh0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* tanh_redundant6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, tanh0)); - HloInstruction* tanh1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, negate0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* tanh2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, tanh1)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* tanh3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, tanh2)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* tuple = builder.AddInstruction( - HloInstruction::CreateTuple({tanh3, negate3, tanh0})); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, - {p0, tanh0, tanh_redundant0, tanh_redundant1, tanh_redundant2, - tanh_redundant3, tanh_redundant4, tanh_redundant5, tanh_redundant6, - negate0, tanh1, negate1, tanh2, negate2, tanh3, negate3, tuple}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpaceUsingCostAnalysis(module.get()); - - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(auto hlo_live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - - std::vector num_live_buffers_in_alternate_mem( - hlo_live_range->flattened_instruction_sequence().size() + 1, 0); - - // Go through each value and for those that are allocated in the alternate - // memory space, increment (inclusive) num_live_buffers_in_alternate_mem for - // every time step that they are live. - for (const HloValue* value : alias_analysis->dataflow_analysis().values()) { - const Shape& shape = value->shape(); - if (!shape.has_layout() || - shape.layout().memory_space() == kDefaultMemorySpace) { - continue; - } - - HloLiveRange::LiveRangeBounds time_bound = - hlo_live_range->buffer_live_ranges().at(value); - for (int i = time_bound.start; i <= time_bound.end; ++i) { - ++num_live_buffers_in_alternate_mem[i]; - } - } - - // The test memory can at most hold two f32[4,3] buffers at a time. If there - // is more than that, it means we have memory corruption. - for (int i = 0; i < num_live_buffers_in_alternate_mem.size(); ++i) { - EXPECT_LE(num_live_buffers_in_alternate_mem[i], 2); - } -} - -TEST_F(MemorySpaceAssignmentTest, - InputOutputsInAlternateMemShouldntBeAssigned) { - // When input/outputs are marked to be in the alternate memory (e.g. - // go/tpu-fast-mem-inference), do not allocate those and assume they will live - // in the alternate memory for the entire computation. The BufferAssignment - // pass, which is run after this, will allocate those buffers. - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - Shape shape_in_alternate_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - Shape shape_in_default_mem = ShapeUtil::MakeShapeWithDenseLayout( - F32, {2, 3}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kDefaultMemorySpace); - // p0 is in the default memory space. - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - // p1 is in the alternate memory space. - HloInstruction* p1 = builder.AddInstruction( - HloInstruction::CreateParameter(1, shape_in_alternate_mem, "p1")); - HloInstruction* negate0 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); - HloInstruction* negate1 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); - HloInstruction* negate2 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); - HloInstruction* negate3 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); - HloInstruction* negate4 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); - HloInstruction* negate5 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); - HloInstruction* negate6 = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); - HloInstruction* add = builder.AddInstruction(HloInstruction::CreateBinary( - shape_in_alternate_mem, HloOpcode::kAdd, negate6, p1)); - // Index {0} of the root instruction is in the alternate memory space, index - // {1} is in the default memory space. - HloInstruction* tuple = - builder.AddInstruction(HloInstruction::CreateTuple({add, negate5})); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, - {p0, p1, negate0, negate1, negate2, negate3, negate4, - negate5, negate6, add, tuple}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.post_module_scoped_alternate_memory_size_in_bytes = 32; - options.is_allowed_in_alternate_mem_fn = [](const HloValue& value) { - return true; - }; - XLA_VLOG_LINES(3, module->ToString()); - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options)); - // Check that the post-module scoped alternate memory chunk is set correctly. - ASSERT_TRUE(preset_assignments->post_module_scoped_alternate_memory_chunk() - .has_value()); - EXPECT_EQ( - preset_assignments->post_module_scoped_alternate_memory_chunk()->size, - 32); - - XLA_VLOG_LINES(3, module->ToString()); - // Ensure that p1 is in the alternate memory and add, which has p1 as an - // operand, has a direct dependency to p1 (no CopyStart/CopyDone). - EXPECT_THAT(p1, op::ShapeWithLayout(shape_in_default_mem)); - EXPECT_THAT(add, op::Add(op::Negate(), op::CopyDone())); - // Make sure add is still in the alternate memory space. - EXPECT_THAT(add, op::ShapeWithLayout(shape_in_alternate_mem)); - - // Check the preset assignments and ensure the inputs/outputs in the alternate - // memory space aren't in the preset assignments. Inputs/outputs in the - // alternate memory space are left to BufferAssignment to be allocated. - for (const auto& position_and_chunk : preset_assignments->chunks()) { - const HloPosition& position = position_and_chunk.first; - XLA_VLOG_LINES(3, position.instruction->ToString()); - EXPECT_NE(position.instruction, p1); - EXPECT_NE(position.instruction, add); - } -} - -TEST_F(MemorySpaceAssignmentTest, PendingChunkMemoryCorruptionBug) { - // Tests a memory corruption bug where the allocated chunk overlaps with a - // pending chunk. To test this, we provide a new buffer interval compare where - // we prioritize the allocation of sine, cosine, and tanh to create the - // situation: - // - // Max memory - // ------------------------------------------- - // +------------+ - // | b | - // +------------+ - // +-------+ - // | | - // | | - // | a | - // | | +------------+ - // | | | n | - // +-------+ +------------+ - // ------------------------------------------- - // Min memory time -> - // - // - // Then allocating for buffer d, we have these two prefetch buffers - // overlapping: - // - // Max memory - // ------------------------------------------- - // +------------+ +----------+ - // | b | | prefetch | - // +------------+ | for o | - // +-------+ +---------+ | - // | | | | | | - // | | | | | | - // | a | | +----|-----+ - // | | | prefetch| +------------+ - // | | | for m | | n | - // +-------+ +---------+ +------------+ - // ------------------------------------------- - // Min memory time -> - // - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY %Entry { - %param0 = f32[8,3] parameter(0) - %param1 = f32[2,4] parameter(1) - %a = f32[8,3] sine(%param0) - %b = f32[2,4] cosine(%param1) - %d = f32[8,3] tanh(%a) - %c = f32[8,3] negate(%a) - %e = f32[2,4] negate(%b) - %f = f32[2,4] negate(%e) - %g = f32[2,4] negate(%f) - %h = f32[2,4] negate(%g) - %i = f32[2,4] negate(%h) - %j = f32[2,4] negate(%i) - %k = f32[2,4] negate(%j) - %l = f32[2,4] negate(%k) - %m = f32[8,3] negate(%d) - %n = f32[2,4] sine(%l) - %o = f32[8,3] negate(%d) - %p = f32[2,4] negate(%n) - %q = f32[8,3] negate(%m) - ROOT %tuple = (f32[2,4], f32[8,3], f32[8,3]) tuple(%p, %q, %o) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_opcode_priority = [](const HloOpcode& opcode) { - switch (opcode) { - case HloOpcode::kSin: - return 0; - case HloOpcode::kCos: - return 1; - case HloOpcode::kTanh: - return 2; - default: - return 3; - } - }; - - return get_opcode_priority(a.buffer->defining_instruction()->opcode()) < - get_opcode_priority(b.buffer->defining_instruction()->opcode()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - buffer_interval_compare, &prefetch_interval_picker); -} - -TEST_F(MemorySpaceAssignmentTest, WhileAliasedArgumentRequiredAssignmentBug) { - // Tests an overly pessimistic assertion when the same HloValue is passed - // multiple times to a while HLO. We already handle this case that the two - // arguments must alias and get the same allocation in AllocateSegment so the - // assertion isn't necessary. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - while_condition { - param1 = (f32[2,4], f32[2,4], f32[2,4]) parameter(0) - ROOT cond = pred[] constant(true) - } - - while_body { - param2 = (f32[2,4], f32[2,4], f32[2,4]) parameter(0) - gte2 = f32[2,4] get-tuple-element(param2), index=0 - gte3 = f32[2,4] get-tuple-element(param2), index=1 - gte4 = f32[2,4] get-tuple-element(param2), index=2 - add = f32[2,4] add(gte2, gte3) - ROOT tuple2 = (f32[2,4], f32[2,4], f32[2,4]) tuple(add, gte3, gte4) - } - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - a = f32[2,4] negate(param0) - b = f32[2,4] negate(param0) - tuple = (f32[2,4], f32[2,4], f32[2,4]) tuple(a, b, b) - while = (f32[2,4], f32[2,4], f32[2,4]) while(tuple), condition=while_condition, body=while_body - gte1 = f32[2,4] get-tuple-element(while), index=0 - gte2 = f32[2,4] get-tuple-element(while), index=1 - ROOT root = f32[2,4] add(gte1, gte2) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, DisallowedUseBug) { - // When we have a disallowed use (in this case tanh), we aren't allowed to - // allocate this use in alternate memory. However, if we have another use - // after this on the same buffer (o), this use may refer to "a" instead of the - // evicted value, which is illegal because "a" will be allocated in the - // alternate memory space. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[8,3] parameter(0) - param1 = f32[2,4] parameter(1) - a = f32[8,3] cosine(param0) - b = f32[2,4] negate(param1) - d = f32[8,3] negate(a) - c = f32[2,4] negate(b) - e = f32[2,4] negate(c) - f = f32[8,3] tanh(a) - g = f32[2,4] negate(e) - h = f32[2,4] negate(g) - i = f32[2,4] negate(h) - j = f32[2,4] negate(i) - k = f32[2,4] negate(j) - l = f32[2,4] negate(k) - m = f32[2,4] negate(l) - n = f32[2,4] sine(m) - o = f32[8,3] negate(a) - p = f32[2,4] negate(n) - q = f32[8,3] add(o, f) - r = f32[8,3] add(q, d) - ROOT tuple = (f32[2,4], f32[8,3]) tuple(p, r) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_opcode_priority = [](const HloOpcode& opcode) { - switch (opcode) { - case HloOpcode::kSin: - return 0; - case HloOpcode::kCos: - return 1; - case HloOpcode::kTanh: - return 2; - default: - return 3; - } - }; - - return get_opcode_priority(a.buffer->defining_instruction()->opcode()) < - get_opcode_priority(b.buffer->defining_instruction()->opcode()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - Options options = DefaultMemorySpaceOptions(); - options.is_use_allowed_in_alternate_mem_fn = [](const HloUse& use) { - return use.instruction->opcode() != HloOpcode::kTanh; - }; - AssignMemorySpace(module.get(), std::move(options), buffer_interval_compare, - &prefetch_interval_picker); -} - -TEST_F(MemorySpaceAssignmentTest, DisallowedUseBugInWhile) { - // Test for situations where we disallow a use (tanh in this case) in the - // alternate memory space and there is a subsequent use that also requires the - // buffer to be in the default memory space. In this case, the allocation in - // the default memory space might not be the very last one, so we need to - // search the allocation sequence and find the one in the default memory - // space. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=3 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = f32[3]{0} get-tuple-element(p0), index=2 - gte3 = pred[] get-tuple-element(p0), index=3 - add = f32[3]{0} add(gte0, gte0) - negate0 = f32[3]{0} negate(add) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - negate15 = f32[3]{0} negate(gte2) - tanh = f32[3]{0} tanh(gte2) - ROOT tuple = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) tuple(negate14, tanh, gte2, gte3) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy0 = f32[3]{0} copy(p0) - copy1 = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) tuple(copy0, copy0, copy1, p1) - while = (f32[3]{0}, f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - ROOT gte = f32[3]{0} get-tuple-element(while), index=2 - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.is_use_allowed_in_alternate_mem_fn = [](const HloUse& use) { - return use.instruction->opcode() != HloOpcode::kTanh; - }; - AssignMemorySpace(module.get(), std::move(options)); -} - -TEST_F(MemorySpaceAssignmentTest, TwoLiveAllocationValuesBase) { - // In this example, we have enough space to give negate.0 alternate memory, - // and we put put negate.0 at the top of MSA's sort order. So, we expect that - // it will get alternate memory. - // - // We are testing a fix for dual live AllocationsValues, with the following - // setup: - // - HloValue H containing the following positions: negate.0, cp-start.0{0} - // - AllocationValue A0 defined at negate.0 - // - Segment A0.S0 define during [negate.0, cp-start.0] - //. - Segment A0.S1 defined during [cp-start.0, add.0] - // - AllocationValue A1 defined at cp-start.0{0} - // - Segment A1.S0 defined during [cp-start.0, cp-done.0] - // - // A0 and A1 are both live for more than 1 instruction. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - /*00:*/ p.0 = f32[10,10,10,10] parameter(0) - /*01:*/ p.1 = f32[10,10,10,10] parameter(1) - /*02:*/ v.0 = f32[10,10,10,10] add(p.1, p.1) - /*03:*/ negate.0 = f32[10,10,10,10] negate(p.0) - /*04:*/ cp-start.0 = (f32[10,10,10,10], f32[10,10,10,10], u32[], u32[]) collective-permute-start(negate.0), source_target_pairs={{0,1},{2,3}} - /*05:*/ v.1 = f32[10,10,10,10] add(v.0, v.0) - /*06:*/ add.0 = f32[10,10,10,10] add(negate.0, negate.0) - /*07:*/ v.2 = f32[10,10,10,10] add(v.1, v.1) - /*08:*/ cp-done.0 = f32[10,10,10,10] collective-permute-done(cp-start.0) - /*09:*/ v.3 = f32[10,10,10,10] add(v.2, v.2) - /*10:*/ ROOT tuple.0 = (f32[10,10,10,10], f32[10,10,10,10], f32[10,10,10,10]) tuple(add.0, cp-done.0, v.3) - })hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 4 * 10 * 10 * 10 * 10; - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - MsaBufferIntervalCompare buffer_interval_compare = - CreateBufferIntervalCompareFnFromInstructionNames({"negate.0"}); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(1, 10); - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Module after MSA:\n" << module->ToString(); - - HloInstruction* copy0 = FindInstruction(module.get(), "negate.0"); - ASSERT_NE(copy0, nullptr); - EXPECT_EQ(copy0->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - TwoLiveAllocationValuesTwoInstructionOverlap) { - // In this example, we have enough space to give negate.0 alternate memory, - // and we put put negate.0 at the top of MSA's sort order. So, we expect that - // it will get alternate memory. - // - // We are testing a fix for dual live AllocationValues, with the following - // setup: - // - HloValue H containing the following positions: negate.0, cp-start.0{0} - // - AllocationValue A0 defined at negate.0 - // - Segment A0.S0 define during [negate.0, cp-start.0] - //. - Segment A0.S1 defined during [cp-start.0, add.0] - // - AllocationValue A1 defined at cp-start.0{0} - // - Segment A1.S0 defined during [cp-start.0, cp-done.0] - // - // A0 and A1 are both live for 2 instructions - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - /*00:*/ p.0 = f32[10,10,10,10] parameter(0) - /*01:*/ p.1 = f32[10,10,10,10] parameter(1) - /*02:*/ v.0 = f32[10,10,10,10] add(p.1, p.1) - /*03:*/ negate.0 = f32[10,10,10,10] negate(p.0) - /*04:*/ cp-start.0 = (f32[10,10,10,10], f32[10,10,10,10], u32[], u32[]) collective-permute-start(negate.0), source_target_pairs={{0,1},{2,3}} - /*05:*/ add.0 = f32[10,10,10,10] add(negate.0, negate.0) - /*06:*/ v.1 = f32[10,10,10,10] add(v.0, v.0) - /*07:*/ v.2 = f32[10,10,10,10] add(v.1, v.1) - /*08:*/ cp-done.0 = f32[10,10,10,10] collective-permute-done(cp-start.0) - /*09:*/ v.3 = f32[10,10,10,10] add(v.2, v.2) - /*10:*/ ROOT tuple.0 = (f32[10,10,10,10], f32[10,10,10,10], f32[10,10,10,10]) tuple(add.0, cp-done.0, v.3) - })hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 4 * 10 * 10 * 10 * 10; - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - MsaBufferIntervalCompare buffer_interval_compare = - CreateBufferIntervalCompareFnFromInstructionNames({"negate.0"}); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(1, 10); - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Module after MSA:\n" << module->ToString(); - - HloInstruction* copy0 = FindInstruction(module.get(), "negate.0"); - ASSERT_NE(copy0, nullptr); - EXPECT_EQ(copy0->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - TwoLiveAllocationValuesFirstLiveAllocationValueOutlastsSecond) { - // In this example, we have enough space to give negate.0 alternate memory, - // and we put put negate.0 at the top of MSA's sort order. So, we expect that - // it will get alternate memory. - // - // We are testing a fix for dual live AllocationValues, with the following - // setup: - // - HloValue H containing the following positions: negate.0, cp-start.0{0} - // - AllocationValue A0 defined at negate.0 - // - Segment A0.S0 define during [negate.0, cp-start.0] - //. - Segment A0.S1 defined during [cp-start.0, add.0] - // - Segment A0.S2 defined during [add.0, add.1] - // - AllocationValue A1 defined at cp-start.0{0} - // - Segment A1.S0 defined during [cp-start.0, cp-done.0] - // - // A0 and A1 are both live for more than 1 instruction. A0 is live beyond the - // end of A1. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - /*00:*/ p.0 = f32[10,10,10,10] parameter(0) - /*01:*/ p.1 = f32[10,10,10,10] parameter(1) - /*02:*/ v.0 = f32[10,10,10,10] add(p.1, p.1) - /*03:*/ negate.0 = f32[10,10,10,10] negate(p.0) - /*04:*/ cp-start.0 = (f32[10,10,10,10], f32[10,10,10,10], u32[], u32[]) collective-permute-start(negate.0), source_target_pairs={{0,1},{2,3}} - /*05:*/ v.1 = f32[10,10,10,10] add(v.0, v.0) - /*06:*/ add.0 = f32[10,10,10,10] add(negate.0, negate.0) - /*07:*/ v.2 = f32[10,10,10,10] add(v.1, v.1) - /*08:*/ cp-done.0 = f32[10,10,10,10] collective-permute-done(cp-start.0) - /*09:*/ v.3 = f32[10,10,10,10] add(v.2, v.2) - /*10:*/ add.1 = f32[10,10,10,10] add(add.0, negate.0) - /*11:*/ ROOT tuple.0 = (f32[10,10,10,10], f32[10,10,10,10], f32[10,10,10,10]) tuple(add.1, cp-done.0, v.3) - })hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 4 * 10 * 10 * 10 * 10; - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - MsaBufferIntervalCompare buffer_interval_compare = - CreateBufferIntervalCompareFnFromInstructionNames({"negate.0"}); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(1, 10); - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Module after MSA:\n" << module->ToString(); - - HloInstruction* copy0 = FindInstruction(module.get(), "negate.0"); - ASSERT_NE(copy0, nullptr); - EXPECT_EQ(copy0->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - TwoLiveAllocationValuesUnableToAllocateContiguousAltMem) { - // In this example, we have enough space to give v.2 alternate memory, - // and we put v.2 at the top of MSA's sort order. So, we expect that - // it will get alternate memory. Second, we try to give negate.0 alternate - // memory, but we can't. In order to give negate.0 alternate memory, we need - // to give it contiguous alternate memory during cp-start.0 to cp-done.0. - // (negate.0 and cp-start.0 {0} alias.) However, v.2 is taking too much - // alternate memory to accommodate that request. - // - // We are testing a fix for dual live AllocationValues, with the following - // setup: - // - HloValue H containing the following positions: negate.0, cp-start.0{0} - // - AllocationValue A0 defined at negate.0 - // - Segment A0.S0 define during [negate.0, cp-start.0] - //. - Segment A0.S1 defined during [cp-start.0, add.0] - // - AllocationValue A1 defined at cp-start.0{0} - // - Segment A1.S0 defined during [cp-start.0, cp-done.0] - // - // A0 and A1 are both live for more than 1 instruction. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - /*00:*/ p.0 = f32[10,10,10,10] parameter(0) - /*01:*/ p.1 = f32[10,10,10,10] parameter(1) - /*02:*/ v.0 = f32[10,10,10,10] add(p.1, p.1) - /*03:*/ negate.0 = f32[10,10,10,10] negate(p.0) - /*04:*/ cp-start.0 = (f32[10,10,10,10], f32[10,10,10,10], u32[], u32[]) collective-permute-start(negate.0), source_target_pairs={{0,1},{2,3}} - /*05:*/ v.1 = f32[10,10,10,10] add(v.0, v.0) - /*06:*/ add.0 = f32[10,10,10,10] add(negate.0, negate.0) - /*07:*/ v.2 = f32[10,10,10,10] add(v.1, v.1) - /*08:*/ cp-done.0 = f32[10,10,10,10] collective-permute-done(cp-start.0) - /*09:*/ v.3 = f32[10,10,10,10] add(v.2, v.2) - /*10:*/ ROOT tuple.0 = (f32[10,10,10,10], f32[10,10,10,10], f32[10,10,10,10]) tuple(add.0, cp-done.0, v.3) - })hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 4 * 10 * 10 * 10 * 10; - MsaBufferIntervalCompare buffer_interval_compare = - CreateBufferIntervalCompareFnFromInstructionNames({"v.2", "negate.0"}); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(1, 10); - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Module after MSA:\n" << module->ToString(); - - HloInstruction* v2 = FindInstruction(module.get(), "v.2"); - ASSERT_NE(v2, nullptr); - EXPECT_EQ(v2->shape().layout().memory_space(), kAlternateMemorySpace); - HloInstruction* copy0 = FindInstruction(module.get(), "negate.0"); - ASSERT_NE(copy0, nullptr); - EXPECT_NE(copy0->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, AvoidRedundantEvictionInWhile) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - add = f32[3]{0} add(negate14, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - gte = f32[3]{0} get-tuple-element(while), index=1 - ROOT negate = f32[3]{0} negate(gte) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect that while{1} is allocated to alternate memory space. Also expect - // that this buffer is prefetched at the end of the while loop body but is - // never evicted (since it has a copy in the default memory space). - const HloInstruction* while_instr = FindInstruction(module.get(), "while"); - EXPECT_EQ(while_instr->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - const HloInstruction* gte1 = FindInstruction(module.get(), "gte1"); - EXPECT_EQ(gte1->user_count(), 1); - EXPECT_EQ(gte1->users()[0]->opcode(), HloOpcode::kTanh); - const HloInstruction* while_root = - while_instr->while_body()->root_instruction(); - EXPECT_THAT(while_root->operand(1), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::Parameter(0)))); -} - -TEST_F(MemorySpaceAssignmentTest, - RedundantEvictionEliminationShouldntAddRedundantParam) { - // Check that if there wasn't an eviction in the while loop, we don't add the - // buffer in default memory as an additional parameter to the while loop. - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - add = f32[3]{0} add(negate1, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - gte = f32[3]{0} get-tuple-element(while), index=1 - ROOT negate = f32[3]{0} negate(gte) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect that while tuple shape contains 3 elements like the original. - const HloInstruction* while_instr = FindInstruction(module.get(), "while"); - EXPECT_EQ(while_instr->shape().tuple_shapes().size(), 3); -} - -TEST_F(MemorySpaceAssignmentTest, AvoidRedundantEvictionInNestedWhile) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond2 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body2 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - add = f32[3]{0} add(negate14, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - while_cond1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT while2 = (f32[3]{0}, f32[3]{0}, pred[]) while(p0), condition=while_cond2, body=while_body2 - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while1 = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond1, body=while_body1 - gte = f32[3]{0} get-tuple-element(while1), index=1 - ROOT negate = f32[3]{0} negate(gte) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect that while1{1} and while2{1} are allocated to alternate memory - // space. Also expect that this buffer is prefetched at the end of the while - // loop body but is never evicted (since it has a copy in the default memory - // space). - const HloInstruction* while1_instr = FindInstruction(module.get(), "while1"); - EXPECT_EQ(while1_instr->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - const HloInstruction* while2_instr = FindInstruction(module.get(), "while2"); - EXPECT_EQ(while2_instr->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - const HloInstruction* gte1 = FindInstruction(module.get(), "gte1"); - EXPECT_EQ(gte1->user_count(), 1); - EXPECT_EQ(gte1->users()[0]->opcode(), HloOpcode::kTanh); - const HloInstruction* while_root = - while2_instr->while_body()->root_instruction(); - EXPECT_THAT(while_root->operand(1), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::GetTupleElement(op::Parameter(0)))); -} - -TEST_F(MemorySpaceAssignmentTest, RedundantEvictionEliminationBug) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - add0 = f32[3]{0} add(negate14, tanh) - add1 = f32[3]{0} add(add0, gte1) - negate = f32[3]{0} negate(add1) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add1, negate, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - gte = f32[3]{0} get-tuple-element(while), index=1 - ROOT negate = f32[3]{0} negate(gte) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect that redundant eviction elimination doesn't kick in because - // while{1} is updated within the body. - const HloInstruction* while_instr = FindInstruction(module.get(), "while"); - EXPECT_EQ(while_instr->shape().tuple_shapes().size(), 3); - EXPECT_EQ(while_instr->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - const HloInstruction* gte1 = FindInstruction(module.get(), "gte1"); - EXPECT_EQ(gte1->user_count(), 2); - EXPECT_NE( - absl::c_find_if(gte1->users(), HloPredicateIsOp), - gte1->users().end()); -} - -TEST_F(MemorySpaceAssignmentTest, RedundantEvictionEliminationInChainedWhile) { - // Check against a bug where a while HLO feeding to another while HLO can - // cause a crash if we performa redundant eviction elimination to the first - // while but not the other (while operand/parameter shapes would mismatch). - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - add = f32[3]{0} add(negate14, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - while_cond2 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body2 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - add = f32[3]{0} add(negate0, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while1 = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond1, body=while_body1 - while2 = (f32[3]{0}, f32[3]{0}, pred[]) while(while1), condition=while_cond2, body=while_body2 - gte = f32[3]{0} get-tuple-element(while2), index=1 - ROOT negate = f32[3]{0} negate(gte) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect that while1 has one more value than while2 in its shape. - EXPECT_EQ( - FindInstruction(module.get(), "while1")->shape().tuple_shapes().size(), - FindInstruction(module.get(), "while2")->shape().tuple_shapes().size() + - 1); -} - -TEST_F(MemorySpaceAssignmentTest, AvoidRedundantEvictionAfterWhile) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = pred[] get-tuple-element(p0), index=2 - add = f32[3]{0} add(gte0, gte1) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, add, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - negate0 = f32[3]{0} negate(p0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, negate14, p1) - while = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - gte0 = f32[3]{0} get-tuple-element(while), index=0 - gte1 = f32[3]{0} get-tuple-element(while), index=1 - negate20 = f32[3]{0} negate(gte1) - negate21 = f32[3]{0} negate(negate20) - negate22 = f32[3]{0} negate(negate21) - negate23 = f32[3]{0} negate(negate22) - negate24 = f32[3]{0} negate(negate23) - negate25 = f32[3]{0} negate(negate24) - negate26 = f32[3]{0} negate(negate25) - negate27 = f32[3]{0} negate(negate26) - negate28 = f32[3]{0} negate(negate27) - negate29 = f32[3]{0} negate(negate28) - negate30 = f32[3]{0} negate(negate29) - negate31 = f32[3]{0} negate(negate30) - negate32 = f32[3]{0} negate(negate31) - negate33 = f32[3]{0} negate(negate32) - negate34 = f32[3]{0} negate(negate33) - ROOT add = f32[3]{0} add(negate34, gte0) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - EXPECT_THAT( - module->entry_computation()->root_instruction()->operand(1), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, op::Copy())); -} - -TEST_F(MemorySpaceAssignmentTest, AvoidRedundantEvictionAfterWhile2) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = pred[] get-tuple-element(p0), index=2 - add = f32[3]{0} add(gte0, gte1) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, add, gte2) - } - - while_cond2 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body2 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = pred[] get-tuple-element(p0), index=2 - add = f32[3]{0} add(gte0, gte1) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, add, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - tuple1 = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while1 = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple1), condition=while_cond1, body=while_body1 - gte0 = f32[3]{0} get-tuple-element(while1), index=0 - gte1 = f32[3]{0} get-tuple-element(while1), index=1 - negate0 = f32[3]{0} negate(gte1) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - tuple2 = (f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, negate14, p1) - while2 = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple2), condition=while_cond2, body=while_body2 - gte2 = f32[3]{0} get-tuple-element(while2), index=0 - gte3 = f32[3]{0} get-tuple-element(while2), index=1 - negate20 = f32[3]{0} negate(gte3) - negate21 = f32[3]{0} negate(negate20) - negate22 = f32[3]{0} negate(negate21) - negate23 = f32[3]{0} negate(negate22) - negate24 = f32[3]{0} negate(negate23) - negate25 = f32[3]{0} negate(negate24) - negate26 = f32[3]{0} negate(negate25) - negate27 = f32[3]{0} negate(negate26) - negate28 = f32[3]{0} negate(negate27) - negate29 = f32[3]{0} negate(negate28) - negate30 = f32[3]{0} negate(negate29) - negate31 = f32[3]{0} negate(negate30) - negate32 = f32[3]{0} negate(negate31) - negate33 = f32[3]{0} negate(negate32) - negate34 = f32[3]{0} negate(negate33) - ROOT add = f32[3]{0} add(negate34, gte2) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get()); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - EXPECT_THAT( - module->entry_computation()->root_instruction()->operand(1), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::AsyncCopy(kDefaultMemorySpace, kAlternateMemorySpace, - op::GetTupleElement(op::While())))); -} - -TEST_F(MemorySpaceAssignmentTest, - AfterWhileRedundantEarlierEvictionModifiedBuffer) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - gte2 = pred[] get-tuple-element(p0), index=2 - add = f32[3]{0} add(gte0, gte1) - negate = f32[3]{0} negate(gte0) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(negate, add, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - copy = f32[3]{0} copy(p0) - negate0 = f32[3]{0} negate(p0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, negate14, p1) - while = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple), condition=while_cond, body=while_body - gte0 = f32[3]{0} get-tuple-element(while), index=0 - gte1 = f32[3]{0} get-tuple-element(while), index=1 - negate20 = f32[3]{0} negate(gte1) - negate21 = f32[3]{0} negate(negate20) - negate22 = f32[3]{0} negate(negate21) - negate23 = f32[3]{0} negate(negate22) - negate24 = f32[3]{0} negate(negate23) - negate25 = f32[3]{0} negate(negate24) - negate26 = f32[3]{0} negate(negate25) - negate27 = f32[3]{0} negate(negate26) - negate28 = f32[3]{0} negate(negate27) - negate29 = f32[3]{0} negate(negate28) - negate30 = f32[3]{0} negate(negate29) - negate31 = f32[3]{0} negate(negate30) - negate32 = f32[3]{0} negate(negate31) - negate33 = f32[3]{0} negate(negate32) - negate34 = f32[3]{0} negate(negate33) - ROOT add = f32[3]{0} add(negate34, gte0) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - EXPECT_THAT( - module->entry_computation()->root_instruction()->operand(1), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::AsyncCopy(kDefaultMemorySpace, kAlternateMemorySpace, - op::GetTupleElement(op::While())))); -} - -TEST_F(MemorySpaceAssignmentTest, - WhileRedundantEvictionWithInefficientAllocationBug) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - while_cond { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - add = f32[3]{0} add(negate1, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - while_cond1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - ROOT gte = pred[] get-tuple-element(p0), index=2 - } - - while_body1 { - p0 = (f32[3]{0}, f32[3]{0}, pred[]) parameter(0) - gte0 = f32[3]{0} get-tuple-element(p0), index=0 - gte2 = pred[] get-tuple-element(p0), index=2 - negate0 = f32[3]{0} negate(gte0) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - gte1 = f32[3]{0} get-tuple-element(p0), index=1 - tanh = f32[3]{0} tanh(gte1) - add = f32[3]{0} add(negate14, tanh) - ROOT tuple = (f32[3]{0}, f32[3]{0}, pred[]) tuple(add, gte1, gte2) - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - p2 = f32[3]{0} parameter(2) - copy = f32[3]{0} copy(p0) - tuple1 = (f32[3]{0}, f32[3]{0}, pred[]) tuple(copy, p0, p1) - while1 = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple1), condition=while_cond, body=while_body - gte0 = f32[3]{0} get-tuple-element(while1), index=0 - gte1 = f32[3]{0} get-tuple-element(while1), index=1 - negate0_entry = f32[3]{0} negate(gte1) - gte2 = pred[] get-tuple-element(while1), index=2 - tuple2 = (f32[3]{0}, f32[3]{0}, pred[]) tuple(gte0, gte1, gte2) - while2 = (f32[3]{0}, f32[3]{0}, pred[]) while(tuple2), condition=while_cond1, body=while_body1 - negate1 = f32[3]{0} negate(negate0_entry) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - negate10 = f32[3]{0} negate(negate9) - negate11 = f32[3]{0} negate(negate10) - negate12 = f32[3]{0} negate(negate11) - negate13 = f32[3]{0} negate(negate12) - negate14 = f32[3]{0} negate(negate13) - gte = f32[3]{0} get-tuple-element(while2), index=1 - ROOT add = f32[3]{0} add(gte, negate14) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // Inject GetInefficientAllocationSites to mark negate0_entry use as - // inefficient. This triggers a corner case bug where allocating for while2{1} - // in the retry allocation fails to find the previous required allocation in - // default memory, and creates a new one which is wrong. - bool marked_inefficient = false; - options.get_inefficient_allocation_sites_fn = - [&](absl::Span positions) - -> std::vector> { - if (absl::c_find(positions, - HloPosition{FindInstruction(module.get(), "while1"), - {1}}) != positions.end() && - !marked_inefficient) { - LOG(INFO) << "Marking the use inefficient."; - marked_inefficient = true; - return {HloUse{FindInstruction(module.get(), "negate0_entry"), 0}}; - } - return {}; - }; - AssignMemorySpace(module.get(), std::move(options)); -} - -// The injected hook flags the same site on every retry. Once that site's -// default memory requirement is committed, re-adding it is a no op, so the -// retry cannot change anything. Finish must finalize rather than loop: before -// the fix this test hangs. -TEST_F(MemorySpaceAssignmentTest, InefficientAllocationRetryWithoutProgress) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = f32[3]{0} parameter(1) - negate0 = f32[3]{0} negate(p1) - negate1 = f32[3]{0} negate(negate0) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - ROOT add = f32[3]{0} add(negate6, p0) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // The hook flags the add instruction's use of p0 as inefficient. - options.get_inefficient_allocation_sites_fn = - [&](absl::Span positions) - -> std::vector> { - if (absl::c_find(positions, HloPosition{FindInstruction(module.get(), "p0"), - {}}) != positions.end()) { - return {HloUse{FindInstruction(module.get(), "add"), 1}}; - } - return {}; - }; - AssignMemorySpace(module.get(), std::move(options)); -} - -TEST_F(MemorySpaceAssignmentTest, DisablePrefetch) { - absl::string_view hlo_string = R"hlo( - HloModule module, is_scheduled=true - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = f32[3]{0} parameter(1) - negate1 = f32[3]{0} negate(p1) - negate2 = f32[3]{0} negate(negate1) - negate3 = f32[3]{0} negate(negate2) - negate4 = f32[3]{0} negate(negate3) - negate5 = f32[3]{0} negate(negate4) - negate6 = f32[3]{0} negate(negate5) - negate7 = f32[3]{0} negate(negate6) - negate8 = f32[3]{0} negate(negate7) - negate9 = f32[3]{0} negate(negate8) - ROOT add = f32[3]{0} add(negate9, p0) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options = DefaultMemorySpaceOptions(); - options.max_outstanding_prefetches = 0; - AssignMemorySpace(module.get(), std::move(options)); - - EXPECT_THAT(module->entry_computation()->root_instruction()->operand(1), - op::Parameter()); -} - -TEST_F(MemorySpaceAssignmentTest, BitcastRoot) { - // Tests against a bug where the root of entry computation is a bitcast - // instruction and it ends up getting an allocation in the alternate memory. - absl::string_view hlo_string = R"hlo( -HloModule primitive_computation_gather.4, is_scheduled=true - -%while_body { - %param.1 = (s32[], f32[3,3,3]) parameter(0) - %get-tuple-element.32 = s32[] get-tuple-element(%param.1), index=0 - %copy.6 = s32[] copy(s32[] %get-tuple-element.32) - %constant.8 = s32[] constant(1) - %add = s32[] add(s32[] %copy.6, s32[] %constant.8) - %get-tuple-element.35 = f32[3,3,3] get-tuple-element(%param.1), index=1 - negate = f32[3,3,3] negate(get-tuple-element.35) - ROOT %tuple.10 = (s32[], f32[3,3,3]) tuple(s32[] %add, f32[3,3,3] negate) -} - -%while_cond { - %param.0 = (s32[], f32[3,3,3]) parameter(0) - %get-tuple-element = s32[] get-tuple-element(%param.0), index=0 - %constant.3 = s32[] constant(3) - ROOT %compare = pred[] compare(s32[] %get-tuple-element, s32[] %constant.3), direction=LT -} - -ENTRY %primitive_computation_gather.4 (parameter.1: f32[3,10,5], parameter.2: s32[3,1]) -> f32[3,3,3] { - %constant.1 = s32[] constant(0) - %copy.11 = s32[] copy(s32[] %constant.1) - %constant = f32[] constant(0) - %broadcast = f32[3,3,3] broadcast(f32[] %constant), dimensions={} - %tuple.8 = (s32[], f32[3,3,3]) tuple(s32[] %copy.11, f32[3,3,3] %broadcast) - %while = (s32[], f32[3,3,3]) while(%tuple.8), condition=%while_cond, body=%while_body - %get-tuple-element.7 = f32[3,3,3] get-tuple-element(%while), index=1 - ROOT %bitcast.1 = f32[3,3,3] bitcast(f32[3,3,3] %get-tuple-element.7) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - const HloInstruction* root = module->entry_computation()->root_instruction(); - EXPECT_TRUE(!root->shape().has_layout() || - root->shape().layout().memory_space() == kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, PrecoloredBuffer) { - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[8,3] parameter(0) - param1 = f32[2,4] parameter(1) - a = f32[8,3]{1,0:S(1)} cosine(param0) - b = f32[2,4] negate(param1) - d = f32[8,3] negate(a) - c = f32[2,4] negate(b) - e = f32[2,4] negate(c) - f = f32[8,3] negate(d) - g = f32[2,4] negate(e) - h = f32[2,4] negate(g) - i = f32[2,4] negate(h) - j = f32[2,4] negate(i) - k = f32[2,4] negate(j) - l = f32[2,4] negate(k) - m = f32[2,4] negate(l) - n = f32[2,4] negate(m) - o = f32[8,3] negate(f) - p = f32[2,4] negate(n) - q = f32[8,3] add(f, o) - r = f32[8,3] add(q, a) - ROOT tuple = (f32[2,4], f32[8,3]) tuple(p, r) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_opcode_priority = [](const HloOpcode& opcode) { - switch (opcode) { - case HloOpcode::kNegate: - return 0; - case HloOpcode::kAdd: - return 1; - case HloOpcode::kCos: - return 2; - default: - return 3; - } - }; - - return get_opcode_priority(a.buffer->defining_instruction()->opcode()) < - get_opcode_priority(b.buffer->defining_instruction()->opcode()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - Options options = DefaultMemorySpaceOptions(); - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - - const HloInstruction* r = FindInstruction(module.get(), "r"); - const HloInstruction* d = FindInstruction(module.get(), "d"); - const HloInstruction* a = FindInstruction(module.get(), "a"); - // Make sure the r and d operands aren't prefetched. - EXPECT_EQ(r->operand(1), a); - EXPECT_EQ(d->operand(0), a); - // Make sure they are allocated in the alternate memory. - EXPECT_EQ(a->shape().layout().memory_space(), kAlternateMemorySpace); - // Make sure the a buffer has an entry in the preset assignments. - auto a_entry = std::find_if( - preset_assignments->chunks().begin(), preset_assignments->chunks().end(), - [&](std::pair position_and_chunk) { - return position_and_chunk.first.instruction == a; - }); - EXPECT_NE(a_entry, preset_assignments->chunks().end()); -} - -TEST_F(MemorySpaceAssignmentTest, PrecoloredBufferOOM) { - // Same as above but there are two 96-byte values that are pinned to the - // alternate memory (the size of the alternate memory is 128 bytes), which is - // unsatisfiable. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[8,3] parameter(0) - param1 = f32[2,4] parameter(1) - a = f32[8,3]{1,0:S(1)} cosine(param0) - b = f32[2,4] negate(param1) - d = f32[8,3] negate(a) - c = f32[2,4] negate(b) - e = f32[2,4] negate(c) - f = f32[8,3] negate(d) - g = f32[2,4] negate(e) - h = f32[2,4] negate(g) - i = f32[2,4] negate(h) - j = f32[2,4] negate(i) - k = f32[2,4] negate(j) - l = f32[2,4] negate(k) - m = f32[2,4] negate(l) - n = f32[2,4] negate(m) - o = f32[8,3]{1,0:S(1)} negate(f) - p = f32[2,4] negate(n) - q = f32[8,3] add(f, o) - r = f32[8,3] add(q, a) - ROOT tuple = (f32[2,4], f32[8,3]) tuple(p, r) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_opcode_priority = [](const HloOpcode& opcode) { - switch (opcode) { - case HloOpcode::kNegate: - return 0; - case HloOpcode::kAdd: - return 1; - case HloOpcode::kCos: - return 2; - default: - return 3; - } - }; - - return get_opcode_priority(a.buffer->defining_instruction()->opcode()) < - get_opcode_priority(b.buffer->defining_instruction()->opcode()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - Options options = DefaultMemorySpaceOptions(); - auto status_or = AssignMemorySpaceAndReturnStatus( - module.get(), std::move(options), buffer_interval_compare, - &prefetch_interval_picker); - EXPECT_THAT( - status_or.status(), - absl_testing::StatusIs( - tsl::error::FAILED_PRECONDITION, - ::testing::HasSubstr("requires allocation in the alternate memory, " - "which could not be satisfied"))); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpShortLiveRange) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - negate0 = bf16[4]{0} negate(param) - collective-permute-start = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(negate0), source_target_pairs={{0,1},{1,2},{2,3}} - negate1 = bf16[4]{0} negate(param) - negate2 = bf16[4]{0} negate(negate1) - negate3 = bf16[4]{0} negate(negate2) - collective-permute-done = bf16[4]{0} collective-permute-done(collective-permute-start) - ROOT add = add(collective-permute-done, negate3) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect both the source and destination buffers to get alternate memory - // allocations. - HloInstruction* collective_permute_start = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start"); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kAlternateMemorySpace); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpLongLiveRange) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - negate0 = bf16[4]{0} negate(param) - collective-permute-start = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(negate0), source_target_pairs={{0,1},{1,2},{2,3}} - negate1 = bf16[4]{0} negate(param) - negate2 = bf16[4]{0} negate(negate1) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - negate5 = bf16[4]{0} negate(negate4) - negate6 = bf16[4]{0} negate(negate5) - negate7 = bf16[4]{0} negate(negate6) - negate8 = bf16[4]{0} negate(negate7) - negate9 = bf16[4]{0} negate(negate8) - negate10 = bf16[4]{0} negate(negate9) - negate11 = bf16[4]{0} negate(negate10) - negate12 = bf16[4]{0} negate(negate11) - negate13 = bf16[4]{0} negate(negate12) - collective-permute-done = bf16[4]{0} collective-permute-done(collective-permute-start) - ROOT add = add(collective-permute-done, negate13) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect none of the buffers to get alternate memory allocations because of - // the long live range. - HloInstruction* collective_permute_start = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start"); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kDefaultMemorySpace); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpLongLiveRangeInputBufferConsumer) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - negate0 = bf16[4]{0} negate(param) - collective-permute-start = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(negate0), source_target_pairs={{0,1},{1,2},{2,3}} - negate1 = bf16[4]{0} negate(negate0) - negate2 = bf16[4]{0} negate(negate1) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - negate5 = bf16[4]{0} negate(negate4) - negate6 = bf16[4]{0} negate(negate5) - negate7 = bf16[4]{0} negate(negate6) - negate8 = bf16[4]{0} negate(negate7) - negate9 = bf16[4]{0} negate(negate8) - negate10 = bf16[4]{0} negate(negate9) - negate11 = bf16[4]{0} negate(negate10) - negate12 = bf16[4]{0} negate(negate11) - negate13 = bf16[4]{0} negate(negate12) - collective-permute-done = bf16[4]{0} collective-permute-done(collective-permute-start) - ROOT add = add(collective-permute-done, negate13) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect none of the buffers to get alternate memory allocations because of - // the long live range and because negate0 is also used by negate1. - HloInstruction* collective_permute_start = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start"); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kDefaultMemorySpace); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, InPlaceAsyncCollectivePermute) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - negate0 = bf16[4]{0} negate(param) - negate1 = bf16[4]{0} negate(param) - const0 = s32[] constant(0) - const1 = s32[] constant(1) - tuple0 = (s32[]) tuple(const0) - tuple1 = (s32[]) tuple(const1) - collective-permute-start = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(negate0, negate1, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate2 = bf16[4]{0} negate(param) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - collective-permute-done = bf16[4]{0} collective-permute-done(collective-permute-start) - ROOT add = add(collective-permute-done, negate4) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - // Expect both the source and destination buffers to get alternate memory - // allocations. - HloInstruction* collective_permute_start = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start"); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kAlternateMemorySpace); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, InPlaceAsyncCollectivePermuteSameBuffer) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - negate0 = bf16[4]{0} negate(param) - const0 = s32[] constant(0) - const1 = s32[] constant(1) - tuple0 = (s32[]) tuple(const0) - tuple1 = (s32[]) tuple(const1) - collective-permute-start = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(negate0, negate0, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate2 = bf16[4]{0} negate(param) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - collective-permute-done = bf16[4]{0} collective-permute-done(collective-permute-start) - ROOT add = add(collective-permute-done, negate4) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - // Expect both the source and destination buffers to get alternate memory - // allocations. - HloInstruction* collective_permute_start = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start"); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kAlternateMemorySpace); - EXPECT_TRUE(collective_permute_start->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - InPlaceAsyncCollectivePermuteSameBufferChained) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - negate0 = bf16[4]{0} negate(param) - const0 = s32[] constant(0) - const1 = s32[] constant(1) - tuple0 = (s32[]) tuple(const0) - tuple1 = (s32[]) tuple(const1) - collective-permute-start.1 = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(negate0, negate0, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate2 = bf16[4]{0} negate(param) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - collective-permute-done.1 = bf16[4]{0} collective-permute-done(collective-permute-start.1) - collective-permute-start.2 = (bf16[4]{0}, bf16[4]{0}, u32[], u32[]) collective-permute-start(collective-permute-done.1, collective-permute-done.1, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate5 = bf16[4]{0} negate(negate4) - negate6 = bf16[4]{0} negate(negate5) - negate7 = bf16[4]{0} negate(negate6) - collective-permute-done.2 = bf16[4]{0} collective-permute-done(collective-permute-start.2) - ROOT add = add(collective-permute-done.2, negate7) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - // Expect both the source and destination buffers to get alternate memory - // allocations. - HloInstruction* collective_permute_start_1 = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start.1"); - EXPECT_TRUE(collective_permute_start_1->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kAlternateMemorySpace); - EXPECT_TRUE(collective_permute_start_1->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kAlternateMemorySpace); - HloInstruction* collective_permute_start_2 = - module->entry_computation()->GetInstructionWithName( - "collective-permute-start.2"); - EXPECT_TRUE(collective_permute_start_2->shape() - .tuple_shapes(0) - .layout() - .memory_space() == kAlternateMemorySpace); - EXPECT_TRUE(collective_permute_start_2->shape() - .tuple_shapes(1) - .layout() - .memory_space() == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - TupleInPlaceAsyncCollectivePermuteSameBufferChained) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - param2 = bf16[48]{0} parameter(1) - negate0.1 = bf16[48]{0} negate(param2) - negate0.2 = bf16[48]{0} negate(param2) - const0 = s32[] constant(0) - const1 = s32[] constant(1) - tuple0.0 = (s32[]) tuple(const0) - tuple0 = ((s32[]), (s32[])) tuple(tuple0.0, tuple0.0) - tuple1.0 = (s32[]) tuple(const1) - tuple1 = ((s32[]), (s32[])) tuple(tuple1.0, tuple1.0) - tuple2 = (bf16[48]{0}, bf16[48]{0}) tuple(negate0.1, negate0.2) - collective-permute-start.1 = ((bf16[48]{0}, bf16[48]{0}), (bf16[48]{0}, bf16[48]{0}), u32[], u32[]) collective-permute-start(tuple2, tuple2, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate2 = bf16[4]{0} negate(param) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - collective-permute-done.1 = (bf16[48]{0}, bf16[48]{0}) collective-permute-done(collective-permute-start.1) - collective-permute-start.2 = ((bf16[48]{0}, bf16[48]{0}), (bf16[48]{0}, bf16[48]{0}), u32[], u32[]) collective-permute-start(collective-permute-done.1, collective-permute-done.1, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate5 = bf16[4]{0} negate(negate4) - negate6 = bf16[4]{0} negate(negate5) - negate7 = bf16[4]{0} negate(negate6) - collective-permute-done.2 = (bf16[48]{0}, bf16[48]{0}) collective-permute-done(collective-permute-start.2) - gte = bf16[48]{0} get-tuple-element(collective-permute-done.2), index=0 - ROOT root = (bf16[48]{0}, bf16[4]{0}) tuple(gte, negate7) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstruction* cp_done1 = - FindInstruction(module.get(), "collective-permute-done.1"); - EXPECT_EQ(cp_done1->operand(0)->opcode(), HloOpcode::kCollectivePermuteStart); - const HloInstruction* cp_done2 = - FindInstruction(module.get(), "collective-permute-done.2"); - EXPECT_EQ(cp_done2->operand(0)->opcode(), HloOpcode::kCollectivePermuteStart); -} - -TEST_F(MemorySpaceAssignmentTest, - TupleInPlaceAsyncCollectivePermuteSameBuffer) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - param2 = bf16[48]{0} parameter(1) - negate0.1 = bf16[48]{0} negate(param2) - negate0.2 = bf16[48]{0} negate(param2) - const0 = s32[] constant(0) - const1 = s32[] constant(1) - tuple0.0 = (s32[]) tuple(const0) - tuple0 = ((s32[]), (s32[])) tuple(tuple0.0, tuple0.0) - tuple1.0 = (s32[]) tuple(const1) - tuple1 = ((s32[]), (s32[])) tuple(tuple1.0, tuple1.0) - tuple2 = (bf16[48]{0}, bf16[48]{0}) tuple(negate0.1, negate0.1) - tuple3 = (bf16[48]{0}, bf16[48]{0}) tuple(negate0.2, negate0.2) - collective-permute-start.1 = ((bf16[48]{0}, bf16[48]{0}), (bf16[48]{0}, bf16[48]{0}), u32[], u32[]) collective-permute-start(tuple2, tuple3, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - negate2 = bf16[4]{0} negate(param) - negate3 = bf16[4]{0} negate(negate2) - negate4 = bf16[4]{0} negate(negate3) - collective-permute-done.1 = (bf16[48]{0}, bf16[48]{0}) collective-permute-done(collective-permute-start.1) - gte = bf16[48]{0} get-tuple-element(collective-permute-done.1), index=0 - ROOT root = (bf16[48]{0}, bf16[4]{0}) tuple(gte, negate4) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstruction* cp_done1 = - FindInstruction(module.get(), "collective-permute-done.1"); - EXPECT_EQ(cp_done1->operand(0)->opcode(), HloOpcode::kCollectivePermuteStart); -} - -TEST_F(MemorySpaceAssignmentTest, - TupleInPlaceAsyncCollectivePermuteSameBufferRoot) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param = bf16[4]{0} parameter(0) - param2 = bf16[48]{0} parameter(1) - negate0.1 = bf16[48]{0} negate(param2) - negate0.2 = bf16[48]{0} negate(param2) - const0 = s32[] constant(0) - const1 = s32[] constant(1) - tuple0.0 = (s32[]) tuple(const0) - tuple0 = ((s32[]), (s32[])) tuple(tuple0.0, tuple0.0) - tuple1.0 = (s32[]) tuple(const1) - tuple1 = ((s32[]), (s32[])) tuple(tuple1.0, tuple1.0) - tuple2 = (bf16[48]{0}, bf16[48]{0}) tuple(negate0.1, negate0.1) - tuple3 = (bf16[48]{0}, bf16[48]{0}) tuple(negate0.2, negate0.2) - collective-permute-start.1 = ((bf16[48]{0}, bf16[48]{0}), (bf16[48]{0}, bf16[48]{0}), u32[], u32[]) collective-permute-start(tuple2, tuple3, tuple0, tuple1), source_target_pairs={{0,1},{1,2},{2,3}}, slice_sizes={{1}} - ROOT collective-permute-done.1 = (bf16[48]{0}, bf16[48]{0}) collective-permute-done(collective-permute-start.1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstruction* cp_done1 = - FindInstruction(module.get(), "collective-permute-done.1"); - EXPECT_EQ(cp_done1->operand(0)->opcode(), HloOpcode::kCollectivePermuteStart); - ShapeUtil::ForEachSubshape( - cp_done1->shape(), - [&](const Shape& subshape, const ShapeIndex& /*index*/) { - if (subshape.IsArray() && subshape.has_layout()) { - EXPECT_EQ(subshape.layout().memory_space(), kDefaultMemorySpace); - } - }); -} - -TEST_F(MemorySpaceAssignmentTest, TupleInPlaceAsyncCollectivePermuteRoot) { - absl::string_view hlo_string = R"hlo( - HloModule inplace_collective_permute, is_scheduled=true - - ENTRY %inplace_collective_permute { - %param.0 = u32[8,1,1] parameter(0) - %constant.1000 = u32[] constant(1000) - %broadcast.1 = u32[8,1,1] broadcast(u32[] %constant.1000), dimensions={} - %broadcast.2 = u32[8,1,1] broadcast(u32[] %constant.1000), dimensions={} - %tuple.input = (u32[8,1,1], u32[8,1,1]) tuple(u32[8,1,1] %param.0, u32[8,1,1] %param.0) - %tuple.output = (u32[8,1,1], u32[8,1,1]) tuple(u32[8,1,1] %broadcast.1, u32[8,1,1] %broadcast.2) - %constant.0 = s32[] constant(0) - %constant.1 = s32[] constant(1) - %constant.2 = s32[] constant(2) - %indices.0.0.0 = (s32[], s32[], s32[]) tuple(s32[] %constant.0, s32[] %constant.0, s32[] %constant.0) - %indices.1.0.0 = (s32[], s32[], s32[]) tuple(s32[] %constant.1, s32[] %constant.0, s32[] %constant.0) - %indices.2.0.0 = (s32[], s32[], s32[]) tuple(s32[] %constant.2, s32[] %constant.0, s32[] %constant.0) - %indices.000.100 = ((s32[], s32[], s32[]), (s32[], s32[], s32[])) tuple((s32[], s32[], s32[]) %indices.0.0.0, (s32[], s32[], s32[]) %indices.1.0.0) - %indices.000.200 = ((s32[], s32[], s32[]), (s32[], s32[], s32[])) tuple((s32[], s32[], s32[]) %indices.0.0.0, (s32[], s32[], s32[]) %indices.2.0.0) - %indices.000.0 = ((s32[], s32[], s32[]), (s32[], s32[], s32[])) tuple((s32[], s32[], s32[]) %indices.0.0.0, (s32[], s32[], s32[]) %indices.0.0.0) - %input.indices = (((s32[], s32[], s32[]), (s32[], s32[], s32[])), ((s32[], s32[], s32[]), (s32[], s32[], s32[]))) tuple(((s32[], s32[], s32[]), (s32[], s32[], s32[])) %indices.000.100, ((s32[], s32[], s32[]), (s32[], s32[], s32[])) %indices.000.0) - %output.indices = (((s32[], s32[], s32[]), (s32[], s32[], s32[])), ((s32[], s32[], s32[]), (s32[], s32[], s32[]))) tuple(((s32[], s32[], s32[]), (s32[], s32[], s32[])) %indices.000.100, ((s32[], s32[], s32[]), (s32[], s32[], s32[])) %indices.000.200) - %collective-permute-start = ((u32[8,1,1], u32[8,1,1]), (u32[8,1,1], u32[8,1,1]), u32[], u32[]) collective-permute-start((u32[8,1,1], u32[8,1,1]) %tuple.input, (u32[8,1,1], u32[8,1,1]) %tuple.output, (((s32[], s32[], s32[]), (s32[], s32[], s32[])), ((s32[], s32[], s32[]), (s32[], s32[], s32[]))) %input.indices, (((s32[], s32[], s32[]), (s32[], s32[], s32[])), ((s32[], s32[], s32[]), (s32[], s32[], s32[]))) %output.indices), channel_id=42, source_target_pairs={{0,1},{1,0},{1,0},{0,1}}, slice_sizes={{4},{4},{4},{4}} - ROOT %collective-permute-done = (u32[8,1,1], u32[8,1,1]) collective-permute-done(((u32[8,1,1], u32[8,1,1]), (u32[8,1,1], u32[8,1,1]), u32[], u32[]) %collective-permute-start) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); - - const HloInstruction* cp_done = - FindInstruction(module.get(), "collective-permute-done"); - EXPECT_EQ(cp_done->operand(0)->opcode(), HloOpcode::kCollectivePermuteStart); - ShapeUtil::ForEachSubshape( - cp_done->shape(), - [&](const Shape& subshape, const ShapeIndex& /*index*/) { - if (subshape.IsArray() && subshape.has_layout()) { - EXPECT_EQ(subshape.layout().memory_space(), kDefaultMemorySpace); - } - }); -} - -TEST_F(MemorySpaceAssignmentTest, ReservedScopedMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,4] parameter(0) - a = f32[2,4] negate(param0) - b = f32[2,4] negate(a) - c = f32[2,4] negate(b) - d = f32[2,4] negate(c) - e = f32[2,4] negate(d) - ROOT f = f32[2,4] add(e, b) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // Make instruction c reserve 64 bytes in the alternate memory. This should - // prevent both b and c to put their outputs in the alternate memory. - options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - if (instruction->name() == "c") { - return 100; - } - return 0; - }; - AssignMemorySpace(module.get(), std::move(options)); - auto get_memory_space = [&](absl::string_view instruction_name) { - return module->entry_computation() - ->GetInstructionWithName(instruction_name) - ->shape() - .layout() - .memory_space(); - }; - EXPECT_TRUE(get_memory_space("a") == kAlternateMemorySpace); - EXPECT_TRUE(get_memory_space("b") == kDefaultMemorySpace); - EXPECT_TRUE(get_memory_space("c") == kDefaultMemorySpace); - EXPECT_TRUE(get_memory_space("d") == kAlternateMemorySpace); - EXPECT_TRUE(get_memory_space("e") == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConstantAllocationFar) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,4] parameter(0) - const = f32[2,4] constant({...}) - a = f32[2,4] negate(param0) - b = f32[2,4] negate(a) - c = f32[2,4] negate(b) - d = f32[2,4] negate(c) - e = f32[2,4] negate(d) - ROOT negate = f32[2,4] add(const, e) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - EXPECT_TRUE(module->entry_computation() - ->GetInstructionWithName("const") - ->shape() - .layout() - .memory_space() == kDefaultMemorySpace); - EXPECT_TRUE(module->entry_computation() - ->GetInstructionWithName("negate") - ->operand(0) - ->shape() - .layout() - .memory_space() == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConstantAllocationNear) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,4] parameter(0) - a = f32[2,4] negate(param0) - b = f32[2,4] negate(a) - c = f32[2,4] negate(b) - d = f32[2,4] negate(c) - e = f32[2,4] negate(d) - const = f32[2,4] constant({...}) - ROOT negate = f32[2,4] add(const, e) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - EXPECT_TRUE(module->entry_computation() - ->GetInstructionWithName("const") - ->shape() - .layout() - .memory_space() == kDefaultMemorySpace); - EXPECT_TRUE(module->entry_computation() - ->GetInstructionWithName("negate") - ->operand(0) - ->shape() - .layout() - .memory_space() == kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, Repack) { - // We initially perform the following allocations at these offsets. - // - // Max memory - // ------------------------------------------- - // - // - // - // - // +------------+ - // | b | - // +------------+ - // +-------+ +------------+ - // | a | | n | - // +-------+ +------------+ - // ------------------------------------------- - // Min memory time -> - // - // Next up, we try to allocate the prefetch for m. However due to - // fragmentation, this won't be possible: - // - // Max memory - // ------------------------------------------- - // - // - // - // +---------+ - // +------------+ | - // | b | | | - // +------------+ | - // +-------+ | | +------------+ - // | a | | d | | n | - // +-------+ +---------+ +------------+ - // ------------------------------------------- - // Min memory time -> - // - // We then call repack to repack the existing allocations which allows us to - // allocate the prefetch for m: - // - // Max memory - // ------------------------------------------- - // +---------+ - // | | - // | | - // | | - // +-------+ | | - // | a | | d | - // +-------+ +---------+ - // +------------+ +------------+ - // | b | | n | - // +------------+ +------------+ - // ------------------------------------------- - // Min memory time -> - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[8,3] parameter(0) - param1 = f32[2,4] parameter(1) - a = f32[2,4] sine(param1) - b = f32[2,4] cosine(param1) - c = f32[8,3] negate(param0) - j = f32[2,4] negate(a) - d = f32[8,3] tanh(param0) - k = f32[2,4] negate(j) - l = f32[2,4] add(b, k) - m = f32[8,3] negate(d) - n = f32[2,4] sine(l) - o = f32[8,3] negate(m) - p = f32[2,4] negate(n) - q = f32[8,3] negate(m) - ROOT tuple = (f32[2,4], f32[8,3], f32[8,3]) tuple(p, q, o) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_opcode_priority = [](const HloOpcode& opcode) { - switch (opcode) { - case HloOpcode::kSin: - return 0; - case HloOpcode::kCos: - return 1; - case HloOpcode::kTanh: - return 2; - default: - return 3; - } - }; - - return get_opcode_priority(a.buffer->defining_instruction()->opcode()) < - get_opcode_priority(b.buffer->defining_instruction()->opcode()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - absl::flat_hash_map, int64_t> repack_map; - // Move "a" from offset 0 to 32. - repack_map[{2, 0}] = 32; - // Move "b" from offset 32 to 0. - repack_map[{3, 32}] = 0; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map); - Options options = DefaultMemorySpaceOptions(); - options.max_repacks = 1; - options.repacker = &repacker; - AssignMemorySpace(module.get(), std::move(options), buffer_interval_compare, - &prefetch_interval_picker); - - // If repacking succeeds, we should find the buffer for d in alternate memory. - const HloInstruction* d = - module->entry_computation()->GetInstructionWithName("d"); - EXPECT_EQ(d->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, RepackExportsAliasedOffsets) { - // This test is that we are correctly exporting aliased offsets for repacking. - // In this example, the buffer produced at HLO "a" will be allocated first, - // and will consist of four allocations: - // 1) a produced in the alternate memory (and then evicted to the default - // memory). 2) a prefetched to the alternate memory to be used by q and - // while HLOs. 3) a used within the while loop body. 4) the output of while - // HLO, used by u. - // - // Since a will be allocated first (the test is crafted to prioritize sine - // HLO), all four allocations should get the same (zero) offsets. However, - // while allocations 2, 3, and 4 need to be colocated with each other, - // allocation 1 doesn't need to be colocated with the other three. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - while_condition { - param1 = (f32[2,4], f32[2,4]) parameter(0) - ROOT cond = pred[] constant(true) - } - - while_body { - param2 = (f32[2,4], f32[2,4]) parameter(0) - gte2 = f32[2,4] get-tuple-element(param2), index=0 - gte3 = f32[2,4] get-tuple-element(param2), index=1 - add = f32[2,4] add(gte2, gte3) - ROOT tuple2 = (f32[2,4], f32[2,4]) tuple(add, gte3) - } - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - a = f32[2,4] sine(param0) - b = f32[2,4] negate(a) - c = f32[2,4] negate(b) - d = f32[2,4] negate(c) - e = f32[2,4] negate(d) - f = f32[2,4] negate(e) - g = f32[2,4] negate(f) - h = f32[2,4] negate(g) - i = f32[2,4] negate(h) - j = f32[2,4] negate(i) - k = f32[2,4] negate(j) - l = f32[2,4] negate(k) - m = f32[2,4] negate(l) - n = f32[2,4] negate(m) - o = f32[2,4] negate(n) - p = f32[2,4] negate(o) - q = f32[2,4] add(p, a) - tuple = (f32[2,4], f32[2,4]) tuple(q, a) - while = (f32[2,4], f32[2,4]) while(tuple), condition=while_condition, body=while_body - gte0 = f32[2,4] get-tuple-element(while), index=0 - gte1 = f32[2,4] get-tuple-element(while), index=1 - r = f32[2,4] negate(gte0) - s = f32[2,4] negate(r) - t = f32[2,4] negate(s) - constant = f32[] constant(0) - broadcast = f32[8,4] broadcast(constant), dimensions={} - cos = f32[8,4] cosine(broadcast) - u = f32[2,4] add(t, gte1) - v = f32[2,4] add(u, param0) - w = f32[8,4] negate(cos) - ROOT tuple3 = (f32[2,4], f32[8,4]) tuple(v, w) - } - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_opcode_priority = [](const HloOpcode& opcode) { - switch (opcode) { - case HloOpcode::kSin: - return 0; - case HloOpcode::kCos: - return 1; - case HloOpcode::kTanh: - return 2; - default: - return 3; - } - }; - - return get_opcode_priority(a.buffer->defining_instruction()->opcode()) < - get_opcode_priority(b.buffer->defining_instruction()->opcode()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - absl::flat_hash_map, int64_t> repack_map; - - // Expect that of the four separate allocations for the "a" buffer, the first - // and the next three are in separate colocations. - auto check_fun = [](absl::Span allocations) { - EXPECT_TRUE(allocations.at(0)->GetColocationsCount() == 1 || - allocations.at(0)->GetColocationsCount() == 3); - EXPECT_EQ(allocations.at(1)->GetColocationsCount(), 3); - EXPECT_EQ(allocations.at(2)->GetColocationsCount(), 3); - EXPECT_TRUE(allocations.at(3)->GetColocationsCount() == 1 || - allocations.at(3)->GetColocationsCount() == 3); - }; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, check_fun); - Options options = DefaultMemorySpaceOptions(); - options.max_repacks = 1; - options.repacker = &repacker; - AssignMemorySpace(module.get(), std::move(options), buffer_interval_compare, - &prefetch_interval_picker); -} - -TEST_F(MemorySpaceAssignmentTest, - RepackExportsAliasedOffsetsForReservedScopedMemory) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,4] parameter(0) - a = f32[2,4] negate(param0) - b = f32[2,4] negate(a) - c = f32[2,4] negate(b) - d = f32[2,4] negate(c) - e = f32[2,4] negate(d) - ROOT f = f32[2,4] add(e, b) -} - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_repacks = 1; - // Make two instructions reserve scoped memory. - options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - if (instruction->name() == "c" || instruction->name() == "d") { - return 100; - } - return 0; - }; - - absl::flat_hash_map, int64_t> repack_map; - bool repacker_ran = false; - - // Expect that the first two value to repack has a colocations size of 2, - // corresponding to the scoped allocations. - auto check_fun = [&](absl::Span allocations) { - EXPECT_EQ(allocations.at(0)->GetColocationsCount(), 2); - EXPECT_EQ(allocations.at(1)->GetColocationsCount(), 2); - repacker_ran = true; - }; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, check_fun); - options.repacker = &repacker; - AssignMemorySpace(module.get(), std::move(options)); - EXPECT_TRUE(repacker_ran); -} - -TEST_F(MemorySpaceAssignmentTest, - ReduceReservedScopedAllocationIfOperandInAlternateMemory) { - // This test is designed to test UpdateReservedScopedAllocationSize() in MSA, - // which will invoke reserved_scoped_memory_fn to update scoped allocation - // size. UpdateReservedScopedAllocationSize() should iterate through all - // scheduled instruction and check if either their operands or outputs has - // been assigned in alternate memory. If so, corresponding operand/output will - // be passed to reserved_scoped_memory_fn. We isolate - // UpdateReservedScopedAllocationSize() by constructing a dummy - // reserved_scoped_memory_fn that return +1 when operand set is empty, and - // return +2 when output set is empty, because if either set of an instruction - // is empty, it is gureented that some scoped allocation is required. We use - // +1/+2 to distinguish the correctness of each set. Therefore, after MSA - // pass, for each instruction, there are a few possible outcomes: - // 1. If both operand set and output set are not empty, scoped allocation - // size should be 0, since reserved_scoped_memory_fn will return 0. - // 2. If only operand set is empty, scoped allocation size should be 2, since - // reserved_scoped_memory_fn will return 2. - // 3. If only output set is empty, scoped allocation size should be 1, since - // reserved_scoped_memory_fn will return 1. - // 4. If both sets are empty, scoped allocation size should be 3. - // Initially, UpdateReservedScopedAllocationSize() will only be invoked after - // each MSA repacking, we use a similar test HLO module as used in "Repack" - // test. This test is capable of testing if - // UpdateReservedScopedAllocationSize() can correctly pass operand/output set - // of all instructions to reserved_scoped_memory_fn. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[8,3] parameter(0) - param1 = f32[2,4] parameter(1) - a = f32[2,4] sine(param1) - b = f32[2,4] cosine(param1) - c = f32[8,3] negate(param0) - j = f32[2,4] negate(a) - d = f32[8,3] tanh(param0) - k = f32[2,4] negate(j) - l = f32[2,4] add(b, k) - m = f32[8,3] negate(d) - n = f32[2,4] sine(l) - o = f32[8,3] negate(m) - p = f32[2,4] negate(n) - q = f32[8,3] negate(m) - ROOT tuple = (f32[2,4], f32[8,3], f32[8,3], f32[8,3]) tuple(p, q, o, c) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - absl::flat_hash_map, int64_t> repack_map; - Options options = DefaultMemorySpaceOptions(); - options.max_repacks = 10; - options.repack_after_every_allocation = true; - options.reduce_scoped_memory_limit = true; - options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - int64_t scoped_memory_size = 0; - if (operands_in_alternate_memory.empty()) { - scoped_memory_size += 1; - LOG(INFO) << instruction->name() - << " has no operand in alternate memory"; - } - if (outputs_in_alternate_memory.empty()) { - scoped_memory_size += 2; - LOG(INFO) << instruction->name() - << " has no output in alternate memory"; - } - return scoped_memory_size; - }; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, nullptr); - options.repacker = &repacker; - std::unique_ptr assignments = - AssignMemorySpace(module.get(), std::move(options)); - // This lambda checks if an instruction's operand has been assigned in - // alternate memory. - auto instruction_consumes_assignment_fn = - [&](absl::string_view instruction_name) -> bool { - HloInstruction* instruction = - module->entry_computation()->GetInstructionWithName(instruction_name); - for (auto& pair : assignments->chunks()) { - HloInstruction* consumer = pair.first.instruction; - if (absl::c_any_of(instruction->operands(), - [&](const HloInstruction* operand) { - return operand == consumer; - })) { - return true; - } - } - return false; - }; - // This lambda checks if an instruction's output has been assigned in - // alternate memory. - auto instruction_produces_assignment_fn = - [&](absl::string_view instruction_name) -> bool { - HloInstruction* instruction = - module->entry_computation()->GetInstructionWithName(instruction_name); - for (auto& pair : assignments->chunks()) { - HloInstruction* producer = pair.first.instruction; - if (producer == instruction) { - return true; - } - } - return false; - }; - auto check_reserved_scoped_memory_fn = - [&](absl::string_view instruction_name) -> bool { - int64_t scoped_memory_size = -1; - for (auto& pair : assignments->scoped_allocation_chunks()) { - HloInstruction* instruction = pair.first; - if (instruction->name() == instruction_name) { - scoped_memory_size = pair.second.size; - } - } - if (!instruction_consumes_assignment_fn(instruction_name)) { - scoped_memory_size -= 1; - } - if (!instruction_produces_assignment_fn(instruction_name)) { - scoped_memory_size -= 2; - } - return scoped_memory_size == 0; - }; - for (auto& pair : assignments->assignment_informations()) { - LOG(INFO) << " space: " << pair.first << ", size: " << pair.second.size; - } - for (auto& pair : assignments->scoped_allocation_chunks()) { - HloInstruction* instruction = pair.first; - LOG(INFO) << instruction->name() << ": " << pair.second.size; - } - EXPECT_TRUE(check_reserved_scoped_memory_fn("a")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("b")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("c")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("j")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("d")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("k")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("l")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("m")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("n")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("o")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("p")); - EXPECT_TRUE(check_reserved_scoped_memory_fn("q")); -} - -TEST_F(MemorySpaceAssignmentTest, ScopedAllocationWithDifferentOffset) { - // This is test is designed against a bug when - // allocate_reserved_scoped_memory_at_same_offset to false, repack block of - // scoped allocation has empty colocations. This is resolved by adding each - // block itself as its own collocation. We test this by checking colocation - // sizes during repacking. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[8,3] parameter(0) - param1 = f32[2,4] parameter(1) - a = f32[2,4] sine(param1) - b = f32[2,4] cosine(param1) - c = f32[8,3] negate(param0) - j = f32[2,4] negate(a) - d = f32[8,3] tanh(param0) - k = f32[2,4] negate(j) - l = f32[2,4] add(b, k) - m = f32[8,3] negate(d) - n = f32[2,4] sine(l) - o = f32[8,3] negate(m) - p = f32[2,4] negate(n) - q = f32[8,3] negate(m) - ROOT tuple = (f32[2,4], f32[8,3], f32[8,3]) tuple(p, q, o) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto check_fun = [](absl::Span allocations) { - for (AllocationBlock* block : allocations) { - if (block->inclusive_start_time == block->end_time) { - EXPECT_GT(block->GetColocationsCount(), 0); - } - } - }; - absl::flat_hash_map, int64_t> repack_map; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, check_fun); - Options options = DefaultMemorySpaceOptions(); - // Scoped allocation needs to have non zero limit. - options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - return 1; - }; - options.max_repacks = 1; - options.repacker = &repacker; - options.allocate_reserved_scoped_memory_at_same_offset = false; - AssignMemorySpace(module.get(), std::move(options)); -} - -TEST_F(MemorySpaceAssignmentTest, - ReduceReservedScopedAllocationUpdatesPeakMemoryUsage) { - // This test is designed to test that the peak_memory_usage_ is updated - // correctly after scoped memory allocation is updated. The test HLO module - // has two HLO values: a and b. The size of a is 64, and the size of b is 128. - // We expect both of them to be allocated in alternate memory. - // - // The reduce_scoped_memory_fn will initially report the scoped - // memory allocation of c as 64. So if b is allocated in the alternate memory, - // the peak memory usage at the time of c will be 196 (64 as scoped memory + - // 128 for its operand). This setup uses up all the memory and prevents a from - // being allocated in alternate memory. However, we also set - // reduce_scoped_memory_fn to report c's scoped memory allocation as 0 if its - // operand is prefetched. So, after repacking, the peak memory usage at the - // time of c should be reduced to 128 (0 as scoped memory + 128 for its - // operand). This allows a to be allocated in the alternate memory. - // - // If the peak_memory_usage_ is not updated correctly, we will hit a check - // failure and a will not be able to be allocated in alternate memory. - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - ENTRY Entry { - param0 = f32[] parameter(0) - param1 = f32[16] parameter(1) - - a = f32[16] negate(param1) - b = f32[32] broadcast(param0), dimensions={} - c = f32[32] sine(b) - d = f32[16] negate(a) - ROOT tuple = (f32[16], f32[32]) tuple(d, c) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - absl::flat_hash_map, int64_t> repack_map; - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 128 + 64; - options.max_repacks = 10; - options.repack_after_every_allocation = true; - options.reduce_scoped_memory_limit = true; - options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) - -> int64_t { - if (instruction->opcode() != HloOpcode::kSin) { - return 0; - } - // Reserve 64 bytes as scoped memory, when we initially allocate scoped - // memory for c. After b is prefetched, reduce the scoped memory to 0. - if (operands_in_alternate_memory.empty()) { - return 64; - } - return 0; - }; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, nullptr); - options.repacker = &repacker; - // TODO(b/563782724): MSA emits a zero-sized alternate memory allocation for - // this module, which trips a CHECK in HeapSimulator::Chunk::OverlapsWith - // during verification. Re-enable verification once that is fixed. - options.verify = false; - AssignMemorySpace(module.get(), std::move(options)); -} - -TEST_F(MemorySpaceAssignmentTest, - RepackShouldntEraseRequiredAssignmentForConditionalOutput) { - // This is a test case for b/171040271. Repacks erase the required assignments - // (since some required assignments are inserted conditionally based on - // allocation decisions), including the fact that conditional outputs are - // always required to get assignments in the default memory. After repacking, - // this required assignment was never added back, causing conditionals to get - // alternate-memory allocations. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3]) parameter(0) - gte = f32[3] get-tuple-element(p0), index=0 - neg1 = f32[3] negate(gte) - ROOT tuple1 = (f32[3]) tuple(neg1) - } - - false_computation { - p0 = (f32[3]) parameter(0) - gte = f32[3] get-tuple-element(p0), index=0 - neg2 = f32[3] negate(gte) - ROOT tuple2 = (f32[3]) tuple(neg2) - } - - ENTRY entry { - p0 = f32[3] parameter(0) - p1 = pred[] parameter(1) - copy = f32[3] copy(p0) - tuple = (f32[3]) tuple(copy) - conditional = (f32[3]) conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - ROOT gte = f32[3] get-tuple-element(conditional), index=0 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - absl::flat_hash_map, int64_t> repack_map; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, nullptr, - /*always_return_modified=*/true); - Options options = DefaultMemorySpaceOptions(); - options.max_repacks = 10; - options.repacker = &repacker; - options.repack_after_every_allocation = true; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - AssignMemorySpace(module.get(), std::move(options), - /*buffer_interval_compare=*/{}, &prefetch_interval_picker); -} - -TEST_F(MemorySpaceAssignmentTest, Determinism) { - // Run memory space assignment a few times to make sure every time it compiles - // to the same thing. - std::unique_ptr module = CreateEvictAndPrefetchModule(); - - AssignMemorySpace(module.get()); - std::string module_str = module->ToString(); - - for (int i = 0; i < 10; ++i) { - std::unique_ptr other_module = CreateEvictAndPrefetchModule(); - AssignMemorySpace(other_module.get()); - EXPECT_EQ(module_str, other_module->ToString()); - } -} - -TEST_F(MemorySpaceAssignmentTest, InPlaceOp) { - // Tests that in-place ops like DynamicUpdateSlice get the same allocation as - // its input. - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation { - param0 = f32[2,3] parameter(0) - constant.1 = f32[] constant(0) - broadcast = f32[2,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - ROOT dynamic-update-slice.5 = f32[2,3] dynamic-update-slice(param0, broadcast, constant.3, constant.3) -} - -ENTRY main { - param = f32[2,3] parameter(0) - negate = f32[2,3] negate(param) - fusion = f32[2,3] fusion(negate), kind=kLoop, calls=fused_computation - ROOT add = f32[2,3] add(fusion, fusion) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace(module.get()); - HloInstruction* negate_instruction = - module->entry_computation()->GetInstructionWithName("negate"); - int64_t negate_offset = - GetAlternateMemoryOffset(*preset_assignments, negate_instruction); - HloInstruction* fusion_instruction = - module->entry_computation()->GetInstructionWithName("fusion"); - int64_t fusion_offset = - GetAlternateMemoryOffset(*preset_assignments, fusion_instruction); - // We expect negate and fusion to get the same offsets. - EXPECT_EQ(negate_offset, fusion_offset); - EXPECT_NE(negate_offset, -1); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalInPlaceOp) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation { - param0 = f32[2,3] parameter(0) - constant.1 = f32[] constant(0) - broadcast = f32[2,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - ROOT dynamic-update-slice.5 = f32[2,3] dynamic-update-slice(param0, broadcast, constant.3, constant.3) -} - -true_computation { - p0 = (f32[2,3]) parameter(0) - gte = f32[2,3] get-tuple-element(p0), index=0 - ROOT neg1 = f32[2,3] negate(gte) -} - -false_computation { - p0 = (f32[2,3]) parameter(0) - gte = f32[2,3] get-tuple-element(p0), index=0 - neg2 = f32[2,3] negate(gte) - ROOT fusion = f32[2,3] fusion(neg2), kind=kLoop, calls=fused_computation -} - -ENTRY entry { - p0 = f32[2,3] parameter(0) - p1 = pred[] parameter(1) - copy = f32[2,3] copy(p0) - tuple = (f32[2,3]) tuple(copy) - ROOT conditional = f32[2,3] conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncCallDisableAlternateMem) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -called_comp { - p0 = f32[2,3] parameter(0) - negate10 = f32[2,3] negate(p0) - negate11 = f32[2,3] negate(negate10) - negate12 = f32[2,3] negate(negate11) - negate13 = f32[2,3] negate(negate12) - negate14 = f32[2,3] negate(negate13) - ROOT negate15 = f32[2,3] negate(negate14) -}, execution_thread="foobar" - -async_comp { - p0 = f32[2,3] parameter(0) - ROOT call = f32[2,3] call(p0), to_apply=called_comp -}, execution_thread="foobar" - -ENTRY entry { - p0 = f32[2,3] parameter(0) - negate0 = f32[2,3] negate(p0) - negate1 = f32[2,3] negate(negate0) - negate2 = f32[2,3] negate(negate1) - negate3 = f32[2,3] negate(negate2) - negate4 = f32[2,3] negate(negate3) - async-start = ((f32[2,3]), f32[2,3], f32[2]) async-start(negate1), async_execution_thread="foobar", calls=async_comp - async-done = f32[2,3] async-done(async-start), async_execution_thread="foobar", calls=async_comp - add0 = f32[2,3] add(negate0, async-done) - negate5 = f32[2,3] negate(add0) - negate6 = f32[2,3] negate(negate5) - negate7 = f32[2,3] negate(negate6) - negate8 = f32[2,3] negate(negate7) - negate9 = f32[2,3] negate(negate8) - ROOT add1 = f32[2,3] add(negate9, async-done) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.is_use_allowed_in_alternate_mem_fn = [](const HloUse& use) { - return use.instruction->opcode() != HloOpcode::kAsyncStart && - use.instruction->opcode() != HloOpcode::kAsyncDone && - use.instruction->parent()->IsMainThread(); - }; - options.is_position_allowed_in_alternate_mem_fn = [](const HloPosition& pos) { - return pos.instruction->opcode() != HloOpcode::kAsyncStart && - pos.instruction->opcode() != HloOpcode::kAsyncDone && - pos.instruction->parent()->IsMainThread(); - }; - AssignMemorySpace(module.get(), std::move(options)); - auto has_alternate_memory_allocation = - [&](const HloInstruction* instruction) { - bool result = false; - auto shape_has_alternate_memory_allocation = - [&](const Shape& subshape, const ShapeIndex& /*index*/) { - if (subshape.IsArray() && - subshape.layout().memory_space() == kAlternateMemorySpace) { - result = true; - } - }; - ShapeUtil::ForEachSubshape(instruction->shape(), - shape_has_alternate_memory_allocation); - for (const HloInstruction* operand : instruction->operands()) { - ShapeUtil::ForEachSubshape(operand->shape(), - shape_has_alternate_memory_allocation); - } - return result; - }; - - // Check that the async ops themselves and the instructions inside async - // computations do not have any alternate memory allocations. - const HloInstruction* async_start = - FindInstruction(module.get(), "async-start"); - const HloInstruction* async_done = - FindInstruction(module.get(), "async-done"); - EXPECT_FALSE(has_alternate_memory_allocation(async_start)); - EXPECT_FALSE(has_alternate_memory_allocation(async_done)); - for (const HloInstruction* instruction : - async_start->async_wrapped_instruction() - ->called_computations()[0] - ->instructions()) { - EXPECT_FALSE(has_alternate_memory_allocation(instruction)); - } - // Check that we still allow the tensor used/produced by the async computation - // to be placed in the alternate memory before/after the async computation. - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Add(op::Negate(), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::AsyncDone()))); - EXPECT_THAT(async_start, - op::AsyncStart(op::AsyncCopy( - kDefaultMemorySpace, kAlternateMemorySpace, op::Negate()))); -} - -TEST_F(MemorySpaceAssignmentTest, InefficientAllocation) { - // The DUS in the fusion only accesses 1/3 of its input/output. The fusion - // input/output buffer is a program input/output buffer, so it creates an - // prefetch and an eviction. When we turn on detecting inefficient - // allocations, we should catch this case and allocate the fusion input/output - // in the default memory space. - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation { - param0 = f32[2,3] parameter(0) - constant.1 = f32[] constant(0) - broadcast = f32[2,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - ROOT dynamic-update-slice.5 = f32[2,3] dynamic-update-slice(param0, broadcast, constant.3, constant.3) -} - -ENTRY entry { - p0 = f32[2,3] parameter(0) - p1 = pred[] parameter(1) - p2 = f32[2,3] parameter(2) - neg0 = f32[2,3] negate(p2) - neg1 = f32[2,3] negate(neg0) - neg2 = f32[2,3] negate(neg1) - neg3 = f32[2,3] negate(neg2) - neg4 = f32[2,3] negate(neg3) - neg5 = f32[2,3] negate(neg4) - neg6 = f32[2,3] negate(neg5) - neg7 = f32[2,3] negate(neg6) - fusion = f32[2,3] fusion(p0), kind=kLoop, calls=fused_computation - neg8 = f32[2,3] negate(neg7) - neg9 = f32[2,3] negate(neg8) - neg10 = f32[2,3] negate(neg9) - neg11 = f32[2,3] negate(neg10) - neg12 = f32[2,3] negate(neg11) - neg13 = f32[2,3] negate(neg12) - neg14 = f32[2,3] negate(neg13) - neg15 = f32[2,3] negate(neg14) - ROOT tuple = (f32[2,3], f32[2,3]) tuple(fusion, neg15) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options0 = DefaultMemorySpaceOptions(); - options0.enable_cross_program_prefetch = false; - // Disable inefficiency check. Expect that the fusion output and operand are - // in the alternate memory. - options0.inefficient_use_to_copy_ratio = 0.0; - AssignMemorySpaceUsingCostAnalysis( - module.get(), - /*memory_space_options_override=*/std::move(options0)); - EXPECT_THAT( - module->entry_computation()->root_instruction(), - op::Tuple(op::AsyncCopy(kDefaultMemorySpace, kAlternateMemorySpace, - op::Fusion(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter()))), - op::Negate())); - - // Re-run MSA with inefficient use-to-copy ratio of 0.5. The fusion only uses - // 8B of data (f32[2,1]) but copies 48B of data (prefetch and eviction of - // f32[2,3]), so this should be considered inefficient (8/48 < 0.5). - ASSERT_OK_AND_ASSIGN(module, ParseAndReturnVerifiedModule(hlo_string)); - Options options1 = DefaultMemorySpaceOptions(); - options1.enable_cross_program_prefetch = false; - options1.inefficient_use_to_copy_ratio = 0.5; - AssignMemorySpaceUsingCostAnalysis( - module.get(), - /*memory_space_options_override=*/std::move(options1)); - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Tuple(op::Fusion(op::Parameter()), op::Negate())); -} - -TEST_F(MemorySpaceAssignmentTest, InefficientAllocationLivelockBug) { - // This is a carefully crafted test where two in-place operations on the same - // buffer (fusion.1 and fusion.2) have a single very long executing operation - // between them. This test deliberately sets a very low transcendentals per - // second value to ensure the tanh op takes longer than what is allowed for a - // no-copy allocation. A quirk of the prefetch interval picker allows a - // prefetch to be scheduled during this tanh operation even though a no-copy - // allocation isn't allowed. Because of this, the first time this buffer is - // allocated, fusion.1 will be put in the alternate memory, but not fusion.2 - // because tanh executes for too long to allow a no-copy allocation. Then, we - // check for inefficient allocations, and consider fusion.1 to be inefficient, - // and add a required assignment in default memory for fusion.1 and - // reallocate. When we reallocate, we aren't allowed to prefetch into - // fusion.1, but fusion.2 succeeds. We then find fusion.2 to be inefficient, - // so we throw away the required assignment on fusion.1 and reallocate. - // Without the appropriate fix, this will go on forever, causing a livelock. - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation_1 { - param0 = f32[5,4] parameter(0) - constant.1 = f32[] constant(0) - broadcast = f32[5,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - ROOT dynamic-update-slice.5 = f32[5,4] dynamic-update-slice(param0, broadcast, constant.3, constant.3) -} - -fused_computation_2 { - param0 = f32[5,4] parameter(0) - constant.1 = f32[] constant(0) - broadcast = f32[5,1] broadcast(constant.1), dimensions={} - constant.3 = s32[] constant(0) - ROOT dynamic-update-slice.5 = f32[5,4] dynamic-update-slice(param0, broadcast, constant.3, constant.3) -} - -ENTRY entry { - p0 = f32[5,4] parameter(0) - p1 = pred[] parameter(1) - p2 = f32[2,3] parameter(2) - neg0 = f32[2,3] negate(p2) - neg1 = f32[2,3] negate(neg0) - neg2 = f32[2,3] negate(neg1) - neg3 = f32[2,3] negate(neg2) - neg4 = f32[2,3] negate(neg3) - neg5 = f32[2,3] negate(neg4) - neg6 = f32[2,3] negate(neg5) - neg7 = f32[2,3] negate(neg6) - fusion.1 = f32[5,4] fusion(p0), kind=kLoop, calls=fused_computation_1 - tanh = f32[2,3] tanh(neg7) - fusion.2 = f32[5,4] fusion(fusion.1), kind=kLoop, calls=fused_computation_2 - neg8 = f32[2,3] negate(tanh) - neg9 = f32[2,3] negate(neg8) - neg10 = f32[2,3] negate(neg0) - neg11 = f32[2,3] negate(neg10) - neg12 = f32[2,3] negate(neg11) - ROOT tuple = (f32[5,4], f32[2,3]) tuple(fusion.2, neg12) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options = DefaultMemorySpaceOptions(); - options.enable_cross_program_prefetch = false; - options.inefficient_use_to_copy_ratio = 0.5; - HloCostAnalysis::Options hlo_cost_options = DefaultHloCostAnalysisOptions(); - hlo_cost_options.set_transcendentals_per_second(0.4); - - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/std::move(options), - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/hlo_cost_options); -} - -TEST_F(MemorySpaceAssignmentTest, - CalledComputationInefficientAllocationLiveLockBug) { - // Parameter 2 is passed into the conditional but it is not actually used. A - // bug in inefficient allocation pinned this buffer to the default memory at - // the use time, but we should really use the "corrected" use time which is - // the earliest-scheduled called computation parameter. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation { - p0 = (f32[3], f32[3]) parameter(0) - gte = f32[3] get-tuple-element(p0), index=0 - neg1 = f32[3] negate(gte) - ROOT tuple1 = (f32[3]) tuple(neg1) - } - - false_computation { - p0 = (f32[3], f32[3]) parameter(0) - gte = f32[3] get-tuple-element(p0), index=0 - neg2 = f32[3] negate(gte) - ROOT tuple2 = (f32[3]) tuple(neg2) - } - - ENTRY entry { - p0 = f32[3] parameter(0) - p1 = pred[] parameter(1) - p2 = f32[3] parameter(2) - copy0 = f32[3] copy(p0) - negate0 = f32[3] negate(p0) - negate1 = f32[3] negate(negate0) - negate2 = f32[3] negate(negate1) - negate3 = f32[3] negate(negate2) - negate4 = f32[3] negate(negate3) - negate5 = f32[3] negate(negate4) - negate6 = f32[3] negate(negate5) - negate7 = f32[3] negate(negate6) - negate8 = f32[3] negate(negate7) - tuple = (f32[3], f32[3]) tuple(copy0, p2) - conditional = (f32[3]) conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - gte = f32[3] get-tuple-element(conditional), index=0 - ROOT add = f32[3] add(gte, negate8) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options = DefaultMemorySpaceOptions(); - options.enable_cross_program_prefetch = false; - options.inefficient_use_to_copy_ratio = 0.5; - HloCostAnalysis::Options hlo_cost_options = DefaultHloCostAnalysisOptions(); - hlo_cost_options.set_transcendentals_per_second(0.4); - - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/std::move(options), - /*cost_analysis_options_override=*/std::nullopt, - /*hlo_cost_options_override=*/hlo_cost_options); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpElapsedTime) { - // Test that async ops are treated to take no time. We assume async operations - // are efficiently scheduled. So, in this example, collective-permute-start - // should take zero time, which should be insufficient time to overlap a - // prefetch for negate1's operand. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = bf16[16]{0} parameter(0) - param1 = bf16[4]{0} parameter(1) - collective-permute-start = (bf16[16]{0}, bf16[16]{0}, u32[], u32[]) collective-permute-start(param0), source_target_pairs={{0,1},{1,2},{2,3}} - negate1 = bf16[4]{0} negate(param1) - collective-permute-done = bf16[16]{0} collective-permute-done(collective-permute-start) - ROOT negate2 = bf16[4]{0} negate(negate1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - AssignMemorySpaceUsingCostAnalysis(module.get()); - EXPECT_THAT(FindInstruction(module.get(), "negate1")->operand(0), - op::Parameter(1)); -} - -TEST_F(MemorySpaceAssignmentTest, AliasedOperandBug) { - // Test for a case where two aliased operands into the same instruction - // (param0 and custom_call2) cause a violation of the required assignment. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[4,4]{0,1} parameter(0) - param1 = f32[4]{0} parameter(1) - param2 = f32[4,4]{0,1} parameter(2) - negate0 = f32[4]{0} negate(param1) - negate1 = f32[4]{0} negate(negate0) - negate2 = f32[4]{0} negate(negate1) - negate3 = f32[4]{0} negate(negate2) - negate4 = f32[4]{0} negate(negate3) - negate5 = f32[4]{0} negate(negate4) - custom_call1 = f32[4,4]{0,1} custom-call(param0), custom_call_target="FooBar", output_to_operand_aliasing={{}: (0, {})} - tanh = f32[4,4]{0,1} tanh(param2) - negate6 = f32[4]{0} negate(negate5) - negate7 = f32[4]{0} negate(negate6) - negate8 = f32[4]{0} negate(negate7) - negate9 = f32[4]{0} negate(negate8) - negate10 = f32[4]{0} negate(negate9) - negate11 = f32[4]{0} negate(negate10) - negate12 = f32[4]{0} negate(negate11) - negate13 = f32[4]{0} negate(negate12) - negate14 = f32[4]{0} negate(negate13) - negate15 = f32[4]{0} negate(negate14) - negate16 = f32[4]{0} negate(negate15) - custom_call2 = f32[4,4]{0,1} custom-call(custom_call1), custom_call_target="FooBar", output_to_operand_aliasing={{}: (0, {})} - custom_call3 = f32[4,4]{0,1} custom-call(param0, custom_call2), custom_call_target="FooBar", output_to_operand_aliasing={{}: (0, {})} - ROOT root = f32[4,4]{0,1} add(tanh, custom_call2) -} - )hlo"; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_inst_priority = [](const HloInstruction* instruction) { - if (instruction->name() == "param2") { - return 0; - } - if (instruction->name() == "param0") { - return 1; - } - return 2; - }; - - return get_inst_priority(a.buffer->defining_instruction()) < - get_inst_priority(b.buffer->defining_instruction()); - }; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - Options options = DefaultMemorySpaceOptions(); - AssignMemorySpace(module.get(), std::move(options), buffer_interval_compare, - &prefetch_interval_picker); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpCustomFusionShortLiveRange) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation_start { - param0 = f32[2,1] parameter(0) - negate = f32[2,1] negate(param0) - ROOT custom-call = (f32[2,1], f32[2,1], u32[], u32[]) custom-call(negate), custom_call_target="AsyncOpStart" -} - -fused_computation_update { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = f32[2,1] parameter(2) - param3 = f32[2,1] parameter(3) - param4 = u32[] parameter(4) - param5 = u32[] parameter(5) - add = f32[2,1] add(param0, param1) - negate = f32[2,1] negate(param2) - ROOT tuple = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) tuple(add, param2, param3, negate, param4, param5) -} - -fused_computation_done { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = u32[] parameter(2) - param3 = u32[] parameter(3) - negate = f32[2,1] negate(param0) - ROOT custom-call = f32[2,1] custom-call(param0, param1, negate, param2, param3), custom_call_target="AsyncOpDone" -} - -ENTRY main { - param = f32[2,1] parameter(0) - negate1 = f32[2,1] negate(param) - negate2 = f32[2,1] negate(negate1) - fusion1 = (f32[2,1], f32[2,1], u32[], u32[]) fusion(negate1), kind=kCustom, output_to_operand_aliasing={{0}: (0, {})}, calls=fused_computation_start - negate3 = f32[2,1] negate(negate2) - negate4 = f32[2,1] negate(negate3) - gte0 = f32[2,1] get-tuple-element(fusion1), index=0 - gte1 = f32[2,1] get-tuple-element(fusion1), index=1 - gte2 = u32[] get-tuple-element(fusion1), index=2 - gte3 = u32[] get-tuple-element(fusion1), index=3 - fusion2 = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) fusion(negate4, negate2, gte0, gte1, gte2, gte3), kind=kLoop, output_to_operand_aliasing={{1}: (2, {}), {2}: (3, {}), {3}: (3, {}), {4}: (4, {}), {5}: (5, {})}, calls=fused_computation_update - gte4 = f32[2,1] get-tuple-element(fusion2), index=0 - negate5 = f32[2,1] negate(gte4) - gte5 = f32[2,1] get-tuple-element(fusion2), index=1 - gte6 = f32[2,1] get-tuple-element(fusion2), index=2 - gte7 = u32[] get-tuple-element(fusion2), index=4 - gte8 = u32[] get-tuple-element(fusion2), index=5 - fusion3 = f32[2,1] fusion(gte5, gte6, gte7, gte8), kind=kCustom, output_to_operand_aliasing={{}: (1, {})}, calls=fused_computation_done - ROOT add = f32[2,1] add(negate5, fusion3) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - absl::string_view inst_name = position.instruction->name(); - if (inst_name == "fusion1" || - (inst_name == "fusion2" && position.index != ShapeIndex({0}))) { - return true; - } - return false; - }; - AssignMemorySpace(module.get(), std::move(options)); - - HloInstruction* fusion1 = - module->entry_computation()->GetInstructionWithName("fusion1"); - HloInstruction* fusion2 = - module->entry_computation()->GetInstructionWithName("fusion2"); - HloInstruction* fusion3 = - module->entry_computation()->GetInstructionWithName("fusion3"); - - EXPECT_THAT(fusion2->operand(2), op::GetTupleElement(fusion1, 0)); - EXPECT_THAT(fusion2->operand(3), op::GetTupleElement(fusion1, 1)); - EXPECT_THAT(fusion3->operand(0), op::GetTupleElement(fusion2, 1)); - EXPECT_THAT(fusion3->operand(1), op::GetTupleElement(fusion2, 2)); - EXPECT_THAT(fusion2->operand(2)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT(fusion2->operand(3)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT(fusion3->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT(fusion3->operand(1)->shape().layout().memory_space(), - kAlternateMemorySpace); - // Operand 0 and 1 should get alternate memory allocations and so is the - // output {0}. - EXPECT_THAT(fusion2->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT(fusion2->operand(1)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT( - ShapeUtil::GetSubshape(fusion2->shape(), {0}).layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpCustomFusionLongLiveRange) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation_start { - param0 = f32[2,1] parameter(0) - negate = f32[2,1] negate(param0) - ROOT custom-call = (f32[2,1], f32[2,1], u32[], u32[]) custom-call(negate), custom_call_target="AsyncOpStart" -} - -fused_computation_update { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = f32[2,1] parameter(2) - param3 = f32[2,1] parameter(3) - param4 = u32[] parameter(4) - param5 = u32[] parameter(5) - add = f32[2,1] add(param0, param1) - negate = f32[2,1] negate(param2) - ROOT tuple = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) tuple(add, param2, param3, negate, param4, param5) -} - -fused_computation_done { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = u32[] parameter(2) - param3 = u32[] parameter(3) - negate = f32[2,1] negate(param0) - ROOT custom-call = f32[2,1] custom-call(param0, param1, negate, param2, param3), custom_call_target="AsyncOpDone" -} - -ENTRY main { - param = f32[2,1] parameter(0) - negate1 = f32[2,1] negate(param) - negate2 = f32[2,1] negate(negate1) - fusion1 = (f32[2,1], f32[2,1], u32[], u32[]) fusion(negate1), kind=kCustom, output_to_operand_aliasing={{0}: (0, {})}, calls=fused_computation_start - negate3 = f32[2,1] negate(negate2) - negate4 = f32[2,1] negate(negate3) - negate5 = f32[2,1] negate(negate4) - negate6 = f32[2,1] negate(negate5) - negate7 = f32[2,1] negate(negate6) - negate8 = f32[2,1] negate(negate7) - negate9 = f32[2,1] negate(negate8) - negate10 = f32[2,1] negate(negate9) - negate11 = f32[2,1] negate(negate10) - negate12 = f32[2,1] negate(negate11) - gte0 = f32[2,1] get-tuple-element(fusion1), index=0 - gte1 = f32[2,1] get-tuple-element(fusion1), index=1 - gte2 = u32[] get-tuple-element(fusion1), index=2 - gte3 = u32[] get-tuple-element(fusion1), index=3 - fusion2 = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) fusion(negate12, negate2, gte0, gte1, gte2, gte3), kind=kLoop, output_to_operand_aliasing={{1}: (2, {}), {2}: (3, {}), {3}: (3, {}), {4}: (4, {}), {5}: (5, {})}, calls=fused_computation_update - gte4 = f32[2,1] get-tuple-element(fusion2), index=0 - negate13 = f32[2,1] negate(gte4) - gte5 = f32[2,1] get-tuple-element(fusion2), index=1 - gte6 = f32[2,1] get-tuple-element(fusion2), index=2 - gte7 = u32[] get-tuple-element(fusion2), index=4 - gte8 = u32[] get-tuple-element(fusion2), index=5 - fusion3 = f32[2,1] fusion(gte5, gte6, gte7, gte8), kind=kCustom, output_to_operand_aliasing={{}: (1, {})}, calls=fused_computation_done - ROOT add = f32[2,1] add(negate13, fusion3) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - absl::string_view inst_name = position.instruction->name(); - if (inst_name == "fusion1" || - (inst_name == "fusion2" && position.index != ShapeIndex({0}))) { - return true; - } - return false; - }; - AssignMemorySpace(module.get(), std::move(options)); - - HloInstruction* fusion1 = - module->entry_computation()->GetInstructionWithName("fusion1"); - HloInstruction* fusion2 = - module->entry_computation()->GetInstructionWithName("fusion2"); - HloInstruction* fusion3 = - module->entry_computation()->GetInstructionWithName("fusion3"); - EXPECT_THAT(fusion2->operand(2), op::GetTupleElement(fusion1, 0)); - EXPECT_THAT(fusion2->operand(2)->shape().layout().memory_space(), - kDefaultMemorySpace); - EXPECT_THAT(fusion2->operand(3), op::GetTupleElement(fusion1, 1)); - EXPECT_THAT(fusion2->operand(3)->shape().layout().memory_space(), - kDefaultMemorySpace); - EXPECT_THAT(fusion3->operand(0), op::GetTupleElement(fusion2, 1)); - EXPECT_THAT(fusion3->operand(0)->shape().layout().memory_space(), - kDefaultMemorySpace); - EXPECT_THAT(fusion3->operand(1), op::GetTupleElement(fusion2, 2)); - EXPECT_THAT(fusion3->operand(1)->shape().layout().memory_space(), - kDefaultMemorySpace); - // Operand 0 and 1 should get alternate memory allocations and so is the - // output {0}. - EXPECT_THAT(fusion2->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT(fusion2->operand(1)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_THAT( - ShapeUtil::GetSubshape(fusion2->shape(), {0}).layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncOpCustomFusionMultipleUsers) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -fused_computation_start { - param0 = f32[2,1] parameter(0) - negate = f32[2,1] negate(param0) - ROOT custom-call = (f32[2,1], f32[2,1], u32[], u32[]) custom-call(negate), custom_call_target="AsyncOpStart" -} - -fused_computation_update1 { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = f32[2,1] parameter(2) - param3 = f32[2,1] parameter(3) - param4 = u32[] parameter(4) - param5 = u32[] parameter(5) - add = f32[2,1] add(param0, param1) - negate = f32[2,1] negate(param2) - ROOT tuple = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) tuple(add, param2, param3, negate, param4, param5) -} - -fused_computation_update2 { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = f32[2,1] parameter(2) - param3 = f32[2,1] parameter(3) - param4 = u32[] parameter(4) - param5 = u32[] parameter(5) - add = f32[2,1] add(param0, param1) - negate = f32[2,1] negate(param2) - ROOT tuple = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) tuple(add, param2, param3, negate, param4, param5) -} - -fused_computation_done { - param0 = f32[2,1] parameter(0) - param1 = f32[2,1] parameter(1) - param2 = u32[] parameter(2) - param3 = u32[] parameter(3) - negate = f32[2,1] negate(param0) - ROOT custom-call = f32[2,1] custom-call(param0, param1, negate, param2, param3), custom_call_target="AsyncOpDone" -} - -ENTRY main { - param = f32[2,1] parameter(0) - negate1 = f32[2,1] negate(param) - negate2 = f32[2,1] negate(negate1) - fusion1 = (f32[2,1], f32[2,1], u32[], u32[]) fusion(negate1), kind=kCustom, output_to_operand_aliasing={{0}: (0, {})}, calls=fused_computation_start - negate3 = f32[2,1] negate(negate2) - negate4 = f32[2,1] negate(negate3) - gte0 = f32[2,1] get-tuple-element(fusion1), index=0 - gte1 = f32[2,1] get-tuple-element(fusion1), index=1 - gte2 = u32[] get-tuple-element(fusion1), index=2 - gte3 = u32[] get-tuple-element(fusion1), index=3 - fusion2 = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) fusion(negate4, negate2, gte0, gte1, gte2, gte3), kind=kLoop, output_to_operand_aliasing={{1}: (2, {}), {2}: (3, {}), {3}: (3, {}), {4}: (4, {}), {5}: (5, {})}, calls=fused_computation_update1 - gte4 = f32[2,1] get-tuple-element(fusion2), index=0 - negate5 = f32[2,1] negate(gte4) - negate10 = f32[2,1] negate(negate5) - negate11 = f32[2,1] negate(negate10) - negate12 = f32[2,1] negate(negate11) - negate13 = f32[2,1] negate(negate12) - negate14 = f32[2,1] negate(negate13) - negate15 = f32[2,1] negate(negate14) - negate16 = f32[2,1] negate(negate15) - negate17 = f32[2,1] negate(negate16) - negate18 = f32[2,1] negate(negate17) - negate19 = f32[2,1] negate(negate18) - fusion3 = (f32[2,1], f32[2,1], f32[2,1], f32[2,1], u32[], u32[]) fusion(negate19, negate2, gte0, gte1, gte2, gte3), kind=kLoop, output_to_operand_aliasing={{1}: (2, {}), {2}: (3, {}), {3}: (3, {}), {4}: (4, {}), {5}: (5, {})}, calls=fused_computation_update2 - gte9 = f32[2,1] get-tuple-element(fusion3), index=0 - negate6 = f32[2,1] negate(gte9) - gte5 = f32[2,1] get-tuple-element(fusion3), index=1 - gte6 = f32[2,1] get-tuple-element(fusion3), index=2 - gte7 = u32[] get-tuple-element(fusion3), index=4 - gte8 = u32[] get-tuple-element(fusion3), index=5 - fusion4 = f32[2,1] fusion(gte5, gte6, gte7, gte8), kind=kCustom, output_to_operand_aliasing={{}: (1, {})}, calls=fused_computation_done - ROOT add = f32[2,1] add(negate6, fusion4) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - absl::string_view inst_name = position.instruction->name(); - if (inst_name == "fusion1" || - (inst_name == "fusion2" && position.index != ShapeIndex({0})) || - (inst_name == "fusion3" && position.index != ShapeIndex({0}))) { - return true; - } - return false; - }; - AssignMemorySpace(module.get(), std::move(options)); -} - -// This test seeks to test that MSA will schedule async copy operations with -// schedule_after=-1 at the very beginning of the program. -// -// The machinery for this is a little opaque from the public API, so we attempt -// to get MSA to self-assign an async copies with schedule_after=-1 by -// exploiting how the hidden algorithm works. This is brittle and subject to -// inadvertent breakage in the future. -TEST_F(MemorySpaceAssignmentTest, HoistCopyStart) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY cross_program_prefetch { - p0 = (f32[8,8]{1,0}, f32[8,2]{1,0}) parameter(0) - get-tuple-element.0 = f32[8,8]{1,0} get-tuple-element(p0), index=0 - add.0 = f32[8,8]{1,0} add(get-tuple-element.0, get-tuple-element.0) - get-tuple-element.1 = f32[8,2]{1,0} get-tuple-element(p0), index=1 - dot.0 = f32[8,2]{1,0} dot(add.0, get-tuple-element.1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot.0) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT dot.1 = f32[2,2]{1,0} dot(negate.8, get-tuple-element.1), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.post_module_scoped_alternate_memory_size_in_bytes = 64; - options.enable_cross_program_prefetch = true; - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options)); - - // Ensure that get-tuple-element.1 is chosen for cross-program prefetch. - auto cross_program_prefetches = module->CrossProgramPrefetches(); - ASSERT_EQ(cross_program_prefetches.size(), 1); - ASSERT_EQ(cross_program_prefetches[0].parameter, 0); - ASSERT_EQ(cross_program_prefetches[0].index, ShapeIndex({1})); - - // Check that the preset assignments contains the correct post module scoped - // allocation chunk, in the presence of cross-program prefetch. - ASSERT_TRUE(preset_assignments->post_module_scoped_alternate_memory_chunk() - .has_value()); - EXPECT_EQ( - preset_assignments->post_module_scoped_alternate_memory_chunk()->size, - 64); - - // Check that the async copy-start for get-tuple-element.1 is hoisted - // after MSA (get-tuple-element.1 was initially the third operation of the - // original schedule). - // - // We expect the only instructions before it are declaring parameter(0) and - // get-tuple-element.1. - for (auto* instruction : module->schedule() - .sequence(module->entry_computation()) - .instructions()) { - auto p0 = op::Parameter(0); - auto get_tuple_element_1 = op::GetTupleElement(p0, 1); - auto copy_start = op::CopyStart(get_tuple_element_1); - EXPECT_THAT(instruction, AnyOf(p0, get_tuple_element_1, copy_start)); - if (::testing::Matches(copy_start)(instruction)) { - EXPECT_TRUE(instruction->cross_program_prefetch_index().has_value()); - break; - } - } -} - -// This test verifies that MSA's internal map of instructions and their operands -// in alternate memory is correct in the presenece of a cross-program prefetch. -// Internally, at the end of MSA, it validates that it's internal map is -// correct. So, as long as this test allocates a cross-program-prefetch in -// alternate memory and doesn't crash, it has passed. -TEST_F(MemorySpaceAssignmentTest, - OperandsInAlternateMemoryWithCrossProgramPrefetch) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY cross_program_prefetch { - p0 = (f32[8,8]{1,0}, f32[8,2]{1,0}) parameter(0) - get-tuple-element.0 = f32[8,8]{1,0} get-tuple-element(p0), index=0 - add.0 = f32[8,8]{1,0} add(get-tuple-element.0, get-tuple-element.0) - get-tuple-element.1 = f32[8,2]{1,0} get-tuple-element(p0), index=1 - dot.0 = f32[8,2]{1,0} dot(add.0, get-tuple-element.1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot.0) - negate.2 = f32[8,2]{1,0} negate(negate.1) - ROOT dot.1 = f32[2,2]{1,0} dot(negate.2, get-tuple-element.1), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.enable_cross_program_prefetch = true; - AssignMemorySpace(module.get(), std::move(options)); - - // Ensure that get-tuple-element.1 is chosen for cross-program prefetch. - auto cross_program_prefetches = module->CrossProgramPrefetches(); - ASSERT_EQ(cross_program_prefetches.size(), 1); -} - -// The test verifies the allocation of exposed operand span buffers (window -// prefetch mode set to kAllocationOnly). It ensures that the buffers are -// allocated at the use time and the fusion instruction operands are updated, -// without scheduling prefetches ahead of time. -TEST_F(MemorySpaceAssignmentTest, OpSpanExposureAllocationOnly) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -%fused_computation { - %p0 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(0) - %p1 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(1) - %p2 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(2) - %add0 = bf16[64,8]{1,0:T(8,128)(2,1)} add(%p0, %p1) - ROOT %add1 = bf16[64,8]{1,0:T(8,128)(2,1)} add(%add0, %p2) -} - -entry { - %p0 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(0) - %p1 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(1) - %p2 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(2) - ROOT fusion = bf16[64,8]{1,0:T(8,128)(2,1)} fusion(bf16[64,8]{1,0:T(8,128)(2,1)} %p0, bf16[64,8]{1,0:T(8,128)(2,1)} %p1, bf16[64,8]{1,0:T(8,128)(2,1)} %p2), kind=kLoop, calls=%fused_computation -} - -)hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options = DefaultMemorySpaceOptions(); - options.window_prefetch_mode = WindowPrefetchMode::kAllocationOnly; - options.op_span_size_fn = - [&](HloInstruction* original_hlo, HloInstruction* cloned_hlo, - int64_t operand_index) -> int64_t { return 32; }; - options.window_prefetch_min_span_size = 2; - AssignMemorySpace(module.get(), std::move(options), - /*max_prefetch_interval=*/10, - /*min_prefetch_interval=*/0); - const HloInstruction* fusion = FindInstruction(module.get(), "fusion"); - // The fusion instruction should have 9 operands: the 3 original operands - // plus 3 exposed operand span buffers, plus 3 sync flags. - EXPECT_EQ(fusion->operand_count(), 9); - - // The added operands are GetTupleElements of WindowPrefetch custom calls. - for (int i = 3; i < 9; i++) { - EXPECT_EQ(fusion->operand(i)->opcode(), HloOpcode::kGetTupleElement); - const HloInstruction* window_prefetch = fusion->operand(i)->operand(0); - EXPECT_TRUE(window_prefetch->IsCustomCall("WindowPrefetch")); - } - - VLOG(2) << "module: " << module->ToString(); -} - -// The test verifies the schedule of prefetching window buffers. We use -// reserved_scoped_memory_fn to reserve a large amount of memory for an -// instruction that's not too far from the fusion and verify that a valid -// schedule can still be found, which prefetches the window buffers right -// after that instruction. -TEST_F(MemorySpaceAssignmentTest, WindowPrefetchSchedule) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -%fused_computation { - %p0 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %neg = bf16[64,8]{1,0:T(8,128)(2,1)} negate(%p0) -} - -entry { - t0 = bf16[64,8]{1,0:T(8,128)(2,1)} parameter(0) - t1 = bf16[64,8]{1,0:T(8,128)(2,1)} negate(t0) - t2 = bf16[64,8]{1,0:T(8,128)(2,1)} negate(t1) - t3 = bf16[64,8]{1,0:T(8,128)(2,1)} fusion(bf16[64,8]{1,0:T(8,128)(2,1)} t0), kind=kLoop, calls=%fused_computation - ROOT t4 = bf16[64,8]{1,0:T(8,128)(2,1)} add(t2, t3) -} - -)hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - // Get info about window prefetch buffers, such as which operands they - // correspond to and their sizes. - auto op_span_size_fn = [&](HloInstruction* original_hlo, - HloInstruction* cloned_hlo, - int64_t operand_index) -> int64_t { - if (original_hlo->name() == "t3" && operand_index == 0) { - return 32; - } - return 0; - }; - - // Set the reserved scoped memory of the negate instruction to be 128MB. This - // will prevent the prefetching of window buffers from starting too early. - auto reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - if (instruction->name() == "t1") { - return 128 * 1024 * 1024; - } - return 0; - }; - - Options options = DefaultMemorySpaceOptions(); - options.window_prefetch_mode = WindowPrefetchMode::kPrefetch; - options.op_span_size_fn = op_span_size_fn; - options.window_prefetch_min_span_size = 2; - options.reserved_scoped_memory_fn = reserved_scoped_memory_fn; - AssignMemorySpace(module.get(), std::move(options), - /*max_prefetch_interval=*/10, - /*min_prefetch_interval=*/0); - const HloInstruction* fusion = FindInstruction(module.get(), "t3"); - // Verify that the fusion instruction is window prefetched. If the fusion - // instruction's operand is window prefetched, the fusion instruction should - // have more than one operand. - EXPECT_GT(fusion->operand_count(), 1); -} - -using AsynchronousCopyOrderingTest = ::testing::Test; - -TEST_F(AsynchronousCopyOrderingTest, Simple) { - // Given asynchronous copies like the following, ensure the pipelining order - // is maintained (earlier start time must have earlier end time). - // 3,11 +-------+ OK - // 1,8 +------+ OK - // 5,14 +--------+ OK - // 7,14 +------+ OK - // 2,16 +-------------+ Violate - // 9,12 +--+ Violate - // 6,17 +----------+ Violate - // 5,13 +-------+ OK (same start as 5,14) - // 5,14 +--------+ OK (same as 5,14) - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyOrdering ordering; - EXPECT_FALSE(ordering.ViolatesOrdering(3, 11)); - ordering.AddCopy({3, 11, 1, alternate_mem_space, 0}); - EXPECT_FALSE(ordering.ViolatesOrdering(1, 8)); - ordering.AddCopy({1, 8, 1, alternate_mem_space, 1}); - EXPECT_FALSE(ordering.ViolatesOrdering(5, 14)); - ordering.AddCopy({5, 14, 1, alternate_mem_space, 2}); - EXPECT_FALSE(ordering.ViolatesOrdering(7, 14)); - ordering.AddCopy({7, 14, 1, alternate_mem_space, 3}); - EXPECT_TRUE(ordering.ViolatesOrdering(2, 16)); - EXPECT_TRUE(ordering.ViolatesOrdering(9, 12)); - EXPECT_TRUE(ordering.ViolatesOrdering(6, 17)); - EXPECT_FALSE(ordering.ViolatesOrdering(5, 13)); - ordering.AddCopy({5, 13, 1, alternate_mem_space, 4}); - EXPECT_FALSE(ordering.ViolatesOrdering(5, 14)); - ordering.AddCopy({5, 14, 1, alternate_mem_space, 5}); -} - -TEST_F(AsynchronousCopyOrderingTest, SameInterval) { - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyOrdering ordering; - EXPECT_FALSE(ordering.ViolatesOrdering(1, 5)); - EXPECT_FALSE(ordering.ViolatesOrdering(2, 4)); - ordering.AddCopy({1, 5, 1, alternate_mem_space, 0}); - EXPECT_TRUE(ordering.ViolatesOrdering(2, 4)); - ordering.AddCopy({1, 5, 1, alternate_mem_space, 1}); - EXPECT_TRUE(ordering.ViolatesOrdering(2, 4)); - ordering.AddCopy({1, 5, 1, alternate_mem_space, 2}); - EXPECT_TRUE(ordering.ViolatesOrdering(2, 4)); - ordering.RemoveCopy({1, 5, 1, alternate_mem_space, 1}); - EXPECT_TRUE(ordering.ViolatesOrdering(2, 4)); - ordering.RemoveCopy({1, 5, 1, alternate_mem_space, 2}); - EXPECT_TRUE(ordering.ViolatesOrdering(2, 4)); - ordering.RemoveCopy({1, 5, 1, alternate_mem_space, 0}); - EXPECT_FALSE(ordering.ViolatesOrdering(2, 4)); -} - -using AsynchronousCopyResourceTest = ::testing::Test; - -TEST_F(AsynchronousCopyResourceTest, Simple) { - // time: 0 1 2 3 4 5 6 7 8 9 - // resource: 2 3 1 6 7 1 7 2 2 4 - // -1,3,5 +-----+ OK - // resource: 0 0 1 6 7 1 7 2 2 4 - // 1,4,4 +---+ OK - // resource: 0 0 0 3 7 1 7 2 2 4 - // 5,9,10 +-----+ - // resource: 0 0 0 3 7 1 0 0 1 4 - // 4,9,3 +-------+ Violate - // 4,8,2 +-----+ OK; The 5,9 copy shifts resource to right. - // resource: 0 0 0 3 7 0 0 0 0 4 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {2.0, 3.0, 1.0, 6.0, 7.0, 1.0, 7.0, 2.0, 2.0, 4.0}); - EXPECT_TRUE(resource.HasEnoughResource(-1, 3, 5.0)); - resource.AddCopy({-1, 3, 5.0, alternate_mem_space, 0}); - EXPECT_TRUE(resource.HasEnoughResource(1, 4, 4.0)); - resource.AddCopy({1, 4, 4.0, alternate_mem_space, 1}); - EXPECT_TRUE(resource.HasEnoughResource(5, 9, 10.0)); - resource.AddCopy({5, 9, 10.0, alternate_mem_space, 2}); - EXPECT_FALSE(resource.HasEnoughResource(4, 9, 3.0)); - EXPECT_TRUE(resource.HasEnoughResource(4, 8, 2.0)); - resource.AddCopy({4, 8, 2.0, alternate_mem_space, 3}); -} - -TEST_F(AsynchronousCopyResourceTest, Propagate) { - // time: 0 1 2 3 4 5 6 7 8 9 - // resource: 2 2 2 2 2 2 2 2 2 2 - // 6,10,2 +-----+ OK - // resource: 2 2 2 2 2 2 2 0 2 2 - // 5,9,2 +-----+ OK - // resource: 2 2 2 2 2 2 0 0 2 2 - // 4,8,2 +-----+ OK - // resource: 2 2 2 2 2 0 0 0 2 2 - // 3,7,2 +-----+ OK - // resource: 2 2 2 2 0 0 0 0 2 2 - // 2,6,2 +-----+ OK - // resource: 2 2 2 0 0 0 0 0 2 2 - // 1,5,2 +-----+ OK - // resource: 2 2 0 0 0 0 0 0 2 2 - // 0,4,3 +-----+ OK - // resource: 2 0 0 0 0 0 0 0 1 2 - // 0,4,3 +-----+ OK - // resource: 2 0 0 0 0 0 0 0 0 0 - // 0,4,1 +-----+ Violate - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0}); - EXPECT_TRUE(resource.HasEnoughResource(6, 10, 2.0)); - resource.AddCopy({6, 10, 2.0, alternate_mem_space, 0}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(5, 9, 2.0)); - resource.AddCopy({5, 9, 2.0, alternate_mem_space, 1}); - EXPECT_TRUE(resource.HasEnoughResource(4, 8, 2.0)); - resource.AddCopy({4, 8, 2.0, alternate_mem_space, 2}); - EXPECT_TRUE(resource.HasEnoughResource(3, 7, 2.0)); - resource.AddCopy({3, 7, 2.0, alternate_mem_space, 3}); - EXPECT_TRUE(resource.HasEnoughResource(2, 6, 2.0)); - resource.AddCopy({2, 6, 2.0, alternate_mem_space, 4}); - EXPECT_TRUE(resource.HasEnoughResource(1, 5, 2.0)); - resource.AddCopy({1, 5, 2.0, alternate_mem_space, 5}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(0, 4, 3.0)); - resource.AddCopy({0, 4, 3.0, alternate_mem_space, 6}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(0, 4, 3.0)); - resource.AddCopy({0, 4, 3.0, alternate_mem_space, 7}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0})); - EXPECT_FALSE(resource.HasEnoughResource(0, 4, 1.0)); -} - -TEST_F(AsynchronousCopyResourceTest, CantPropagate) { - // time: 0 1 2 3 4 5 6 7 8 9 - // resource: 2 2 2 2 2 2 2 2 2 2 - // 5,10,2 +-------+ OK - // resource: 2 2 2 2 2 2 0 2 2 2 - // 4,7,2 +---+ OK - // resource: 2 2 2 2 2 0 0 2 2 2 - // 4,8,4 +-----+ OK - // resource: 2 2 2 2 2 0 0 0 0 2 - // 3,6,4 +---+ Violate - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0}); - EXPECT_TRUE(resource.HasEnoughResource(5, 10, 2.0)); - resource.AddCopy({5, 10, 2.0, alternate_mem_space, 0}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 2.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(4, 7, 2.0)); - resource.AddCopy({4, 7, 2.0, alternate_mem_space, 1}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 2.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(4, 8, 4.0)); - resource.AddCopy({4, 8, 4.0, alternate_mem_space, 2}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0, 2.0})); - EXPECT_FALSE(resource.HasEnoughResource(3, 6, 4.0)); -} - -TEST_F(AsynchronousCopyResourceTest, Nested) { - // time: 0 1 2 3 4 - // resource: 2 2 2 2 2 - // 1,3,2 +-+ OK - // resource: 2 2 0 2 2 - // 0,4,4 +-----+ Violate - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource({2.0, 2.0, 2.0, 2.0, 2.0}); - EXPECT_TRUE(resource.HasEnoughResource(1, 3, 2.0)); - resource.AddCopy({1, 3, 2.0, alternate_mem_space, 0}); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 2.0, 2.0})); - EXPECT_FALSE(resource.HasEnoughResource(0, 4, 4.0)); -} - -TEST_F(AsynchronousCopyResourceTest, Remove) { - // time: 0 1 2 3 4 - // resource: 2 2 2 2 2 - // add:2,5,2 +---+ OK - // resource: 2 2 2 0 2 - // add:-1,2,3+---+ OK - // resource: 0 1 2 0 2 - // add:0,4,4 +-----+ OK - // resource: 0 0 0 0 1 - // rem:0,4,4 +-----+ - // resource: 0 1 2 0 2 - // rem:2,5,2 +---+ - // resource: 0 1 2 2 2 - // rem:-1,2,3+---+ - // resource: 2 2 2 2 2 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource({2.0, 2.0, 2.0, 2.0, 2.0}); - AsynchronousCopy copy1{2, 5, 2.0, alternate_mem_space, 0}; - AsynchronousCopy copy2{-1, 2, 3.0, alternate_mem_space, 1}; - AsynchronousCopy copy3{0, 4, 4.0, alternate_mem_space, 2}; - EXPECT_TRUE(resource.HasEnoughResource(2, 5, 2.0)); - resource.AddCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 0.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(-1, 2, 3.0)); - resource.AddCopy(copy2); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 1.0, 2.0, 0.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(0, 4, 4.0)); - resource.AddCopy(copy3); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 0.0, 0.0, 0.0, 1.0})); - resource.RemoveCopy(copy3); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 1.0, 2.0, 0.0, 2.0})); - resource.RemoveCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 1.0, 2.0, 2.0, 2.0})); - resource.RemoveCopy(copy2); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0})); -} - -TEST_F(AsynchronousCopyResourceTest, NestedRemove) { - // time: 0 1 2 3 4 - // resource: 2 2 2 2 2 - // add:1,3,2 +-+ OK - // resource: 2 2 0 2 2 - // add:0,4,4 +-----+ Violate - // rem:1,3,2 +-+ - // resource: 2 2 2 2 2 - // add:0,4,4 +-----+ OK - // resource: 2 0 0 2 2 - // add:1,3,2 +-+ Violate - // rem:0,4,4 +-----+ - // resource: 2 2 2 2 2 - // add:1,3,2 +-+ OK - // resource: 2 2 0 2 2 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource({2.0, 2.0, 2.0, 2.0, 2.0}); - AsynchronousCopy copy1{1, 3, 2.0, alternate_mem_space, 0}; - AsynchronousCopy copy2{0, 4, 4.0, alternate_mem_space, 1}; - EXPECT_TRUE(resource.HasEnoughResource(1, 3, 2.0)); - resource.AddCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 2.0, 2.0})); - EXPECT_FALSE(resource.HasEnoughResource(0, 4, 4.0)); - resource.RemoveCopy(copy1); - auto current_resources = resource.GetCurrentResources(); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(0, 4, 4.0)); - resource.AddCopy(copy2); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 0.0, 0.0, 2.0, 2.0})); - EXPECT_FALSE(resource.HasEnoughResource(1, 3, 2.0)); - resource.RemoveCopy(copy2); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(1, 3, 2.0)); -} - -TEST_F(AsynchronousCopyResourceTest, PropagateRemove) { - // time: 0 1 2 3 4 5 6 7 8 9 - // resource: 2 2 2 2 2 2 2 2 2 2 - // add:6,10,2 +-----+ OK - // resource: 2 2 2 2 2 2 2 0 2 2 - // add:5,9,2 +-----+ OK - // resource: 2 2 2 2 2 2 0 0 2 2 - // add:4,8,2 +-----+ OK - // resource: 2 2 2 2 2 0 0 0 2 2 - // add:3,7,2 +-----+ OK - // resource: 2 2 2 2 0 0 0 0 2 2 - // add:2,6,2 +-----+ OK - // resource: 2 2 2 0 0 0 0 0 2 2 - // add:1,5,2 +-----+ OK - // resource: 2 2 0 0 0 0 0 0 2 2 - // add:0,4,3 +-----+ OK - // resource: 2 0 0 0 0 0 0 0 1 2 - // add:0,5,3 +-------+ OK - // resource: 2 0 0 0 0 0 0 0 0 0 - // rem:0,5,3 +-------+ - // resource: 2 0 0 0 0 0 0 0 1 2 - // rem:0,4,3 +-----+ - // resource: 2 2 0 0 0 0 0 0 2 2 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0}); - EXPECT_TRUE(resource.HasEnoughResource(6, 10, 2.0)); - resource.AddCopy({6, 10, 2.0, alternate_mem_space, 0}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(5, 9, 2.0)); - resource.AddCopy({5, 9, 2.0, alternate_mem_space, 1}); - EXPECT_TRUE(resource.HasEnoughResource(4, 8, 2.0)); - resource.AddCopy({4, 8, 2.0, alternate_mem_space, 2}); - EXPECT_TRUE(resource.HasEnoughResource(3, 7, 2.0)); - resource.AddCopy({3, 7, 2.0, alternate_mem_space, 3}); - EXPECT_TRUE(resource.HasEnoughResource(2, 6, 2.0)); - resource.AddCopy({2, 6, 2.0, alternate_mem_space, 4}); - EXPECT_TRUE(resource.HasEnoughResource(1, 5, 2.0)); - resource.AddCopy({1, 5, 2.0, alternate_mem_space, 5}); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0})); - AsynchronousCopy copy1{0, 4, 3.0, alternate_mem_space, 6}; - EXPECT_TRUE(resource.HasEnoughResource(0, 4, 3.0)); - resource.AddCopy(copy1); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(0, 5, 3.0)); - AsynchronousCopy copy2{0, 5, 3.0, alternate_mem_space, 7}; - resource.AddCopy(copy2); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0})); - resource.RemoveCopy(copy2); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 2.0})); - resource.RemoveCopy(copy1); - EXPECT_EQ( - resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0})); -} - -TEST_F(AsynchronousCopyResourceTest, StartAtZeroAndRemove) { - // time: 0 1 2 3 4 - // resource: 0 0 1 1 2 - // add:0,4,2 +-----+ OK - // resource: 0 0 0 0 2 - // rem:0,4,2 +-----+ - // resource: 0 0 1 1 2 - // add:0,4,2 +-----+ OK - // resource: 0 0 0 0 2 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource({0.0, 0.0, 1.0, 1.0, 2.0}); - AsynchronousCopy copy1{0, 4, 2.0, alternate_mem_space, 0}; - EXPECT_TRUE(resource.HasEnoughResource(0, 4, 2.0)); - resource.AddCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 0.0, 0.0, 0.0, 2.0})); - resource.RemoveCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 0.0, 1.0, 1.0, 2.0})); - resource.AddCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({0.0, 0.0, 0.0, 0.0, 2.0})); -} - -// Below test only works when the resource values are scaled to int64 to avoid -// floating point precision issues. -TEST_F(AsynchronousCopyResourceTest, ConsumeResourceScaledIntegerResource) { - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {5.71429e-10, 8.71333e-09, 8.71333e-09, 1.74267e-08, 1.74267e-08}); - AsynchronousCopy copy1{0, 2, 8.71333e-09, alternate_mem_space, 0}; - EXPECT_TRUE(resource.HasEnoughResource(0, 2, 8.71333e-09)); - resource.AddCopy(copy1); - - AsynchronousCopy copy2{0, 3, 4.35667e-09, alternate_mem_space, 1}; - EXPECT_TRUE(resource.HasEnoughResource(0, 3, 4.35667e-09)); - resource.AddCopy(copy2); - - AsynchronousCopy copy3{2, 4, 4.35667e-09, alternate_mem_space, 2}; - EXPECT_TRUE(resource.HasEnoughResource(2, 4, 4.35667e-09)); - resource.AddCopy(copy3); - - // This call to RemoveCopy should not cause a crash. - resource.RemoveCopy(copy1); -} - -TEST_F(AsynchronousCopyResourceTest, OutOfOrderRemovalSameStartTime) { - // time: 0 1 2 3 4 - // resource: 2 2 2 2 2 - // add:1,3,1 +-+ OK - // resource: 2 2 1 2 2 - // add:1,4,2 +---+ OK - // resource: 2 2 0 1 2 - // rem:1,3,1 +-+ - // resource: 2 2 0 2 2 - // add:1,5,1 +-----+ OK - // resource: 2 2 0 1 2 - // add:1,5,1 +-----+ OK - // resource: 2 2 0 0 2 - // add:1,5,1 +-----+ OK - // resource: 2 2 0 0 1 - // add:1,5,1 +-----+ OK - // resource: 2 2 0 0 0 - // add:1,5,1 +-----+ Violate - // rem:1,4,2 +---+ - // resource: 2 2 0 0 2 - // rem:1,5,1 +-----+ - // resource: 2 2 0 1 2 - // rem:1,5,1 +-----+ - // resource: 2 2 0 2 2 - // rem:1,5,1 +-----+ - // resource: 2 2 1 2 2 - // rem:1,5,1 +-----+ - // resource: 2 2 2 2 2 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource({2.0, 2.0, 2.0, 2.0, 2.0}); - AsynchronousCopy copy1{1, 3, 1.0, alternate_mem_space, 0}; - AsynchronousCopy copy2{1, 4, 2.0, alternate_mem_space, 1}; - EXPECT_TRUE(resource.HasEnoughResource(1, 3, 1.0)); - resource.AddCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 1.0, 2.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(1, 4, 2.0)); - resource.AddCopy(copy2); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 1.0, 2.0})); - resource.RemoveCopy(copy1); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 2.0, 2.0})); - - AsynchronousCopy copy3{1, 5, 1.0, alternate_mem_space, 2}; - AsynchronousCopy copy4{1, 5, 1.0, alternate_mem_space, 3}; - AsynchronousCopy copy5{1, 5, 1.0, alternate_mem_space, 4}; - AsynchronousCopy copy6{1, 5, 1.0, alternate_mem_space, 5}; - EXPECT_TRUE(resource.HasEnoughResource(1, 5, 1.0)); - resource.AddCopy(copy3); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 1.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(1, 5, 1.0)); - resource.AddCopy(copy4); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 2.0})); - EXPECT_TRUE(resource.HasEnoughResource(1, 5, 1.0)); - resource.AddCopy(copy5); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 1.0})); - EXPECT_TRUE(resource.HasEnoughResource(1, 5, 1.0)); - resource.AddCopy(copy6); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 0.0})); - EXPECT_FALSE(resource.HasEnoughResource(1, 5, 1.0)); - - resource.RemoveCopy(copy2); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 0.0, 2.0})); - resource.RemoveCopy(copy3); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 1.0, 2.0})); - resource.RemoveCopy(copy4); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 0.0, 2.0, 2.0})); - resource.RemoveCopy(copy5); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 1.0, 2.0, 2.0})); - resource.RemoveCopy(copy6); - EXPECT_EQ(resource.GetCurrentResources(), - std::vector({2.0, 2.0, 2.0, 2.0, 2.0})); -} - -TEST_F(AsynchronousCopyResourceTest, HasEnoughResourceMultiCheckSuccess) { - // time: 0 1 2 3 4 5 6 7 8 9 - // resource: 2 1 3 6 7 3 7 2 2 4 - // -1,3,5 +-----+ OK - // resource: 0 0 1 6 7 3 7 2 2 4 - // 1,10,4 +---------------+ OK - // resource: 0 0 0 3 7 3 7 2 2 4 - // 0,6,4 +-----------+ - // 4,6,3 +-+ 2 copies OK; The 1,10 copy shifts. - // resource: 0 0 0 0 6 0 7 2 2 4 - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {2.0, 1.0, 3.0, 6.0, 7.0, 3.0, 7.0, 2.0, 2.0, 4.0}); - EXPECT_TRUE(resource.HasEnoughResource(-1, 3, 5.0)); - resource.AddCopy({-1, 3, 5.0, alternate_mem_space, 0}); - EXPECT_TRUE(resource.HasEnoughResource(1, 10, 4.0)); - resource.AddCopy({1, 10, 4.0, alternate_mem_space, 1}); - - LOG(INFO) << "AsynchronousCopyResource after setup:\n" - << resource.Dump(0, 10, alternate_mem_space); - - // We run the check in a loop to demonstrate that it is not modifying the - // underlying data structures. - for (int i = 0; i < 4; ++i) { - EXPECT_TRUE( - resource.HasEnoughResourceMultiCheck({{0, 6, 4.0}, {4, 6, 3.0}})); - } -} - -TEST_F(AsynchronousCopyResourceTest, HasEnoughResourceMultiCheckFailure) { - // time: 0 1 2 3 4 5 6 7 8 9 - // resource: 2 1 3 6 7 3 7 2 2 4 - // -1,3,5 +-----+ OK - // resource: 0 0 1 6 7 3 7 2 2 4 - // 1,10,4 +---------------+ OK - // resource: 0 0 0 3 7 3 7 2 2 4 - // 0,6,4 +-----------+ - // 4,6,4 +-+ Not-OK - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource( - {2.0, 1.0, 3.0, 6.0, 7.0, 3.0, 7.0, 2.0, 2.0, 4.0}); - EXPECT_TRUE(resource.HasEnoughResource(-1, 3, 5.0)); - resource.AddCopy({-1, 3, 5.0, alternate_mem_space, 0}); - EXPECT_TRUE(resource.HasEnoughResource(1, 10, 4.0)); - resource.AddCopy({1, 10, 4.0, alternate_mem_space, 1}); - - LOG(INFO) << "AsynchronousCopyResource after setup:\n" - << resource.Dump(0, 10, alternate_mem_space); - - EXPECT_FALSE( - resource.HasEnoughResourceMultiCheck({{0, 6, 4.0}, {4, 6, 4.0}})); -} - -TEST_F(AsynchronousCopyResourceTest, - HasEnoughResourceMultiCheckRegressionTest) { - auto alternate_mem_space = MemorySpace::kAlternate; - AsynchronousCopyResource resource({/*0:*/ 24.0f, - /*1:*/ 0.0f, - /*2:*/ 6.0f, - /*3:*/ 411.0f, - /*4:*/ 3479.0f, - /*5:*/ 0.0f, - /*6:*/ 0.0f, - /*7:*/ 1537.0f, - /*8:*/ 3095.0f, - /*9:*/ 0.0f, - /*10:*/ 26.7f}); - AsynchronousCopy copy1({1, 8, 170.8f, alternate_mem_space, 1}); - AsynchronousCopy copy2({2, 8, 170.8f, alternate_mem_space, 2}); - resource.AddCopy(copy1); - resource.AddCopy(copy2); - - LOG(INFO) << "AsynchronousCopyResource after setup:\n" - << resource.Dump(0, 11, alternate_mem_space); - // Under the current AsynchronousCopyResource implementation, this - // HasEnoughResource check fails. Although, other designs could rearrange - // resources in a manner that fits the check. - EXPECT_FALSE( - resource.HasEnoughResourceMultiCheck({{0, 4, 170.8}, {1, 4, 170.8}})); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - HloInstruction* lhs = builder.AddInstruction( - HloInstruction::CreateParameter(0, lhs_shape, "lhs")); - HloInstruction* rhs = builder.AddInstruction( - HloInstruction::CreateParameter(1, rhs_shape, "rhs")); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 1); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Dot(op::Parameter(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::Parameter(1)))); -} - -TEST_F(MemorySpaceAssignmentTest, MultiCrossProgramPrefetchTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kFirstOutput = 4; - constexpr int kSecondOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto first_weight_shape = ShapeUtil::MakeShape(F32, {kFeature, kFirstOutput}); - auto second_weight_shape = - ShapeUtil::MakeShape(F32, {kFirstOutput, kSecondOutput}); - auto intermediate_shape = ShapeUtil::MakeShape(F32, {kBatch, kFirstOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kSecondOutput}); - HloInstruction* lhs = builder.AddInstruction( - HloInstruction::CreateParameter(0, lhs_shape, "lhs")); - HloInstruction* first_weight = builder.AddInstruction( - HloInstruction::CreateParameter(1, first_weight_shape, "first_weight")); - HloInstruction* second_weight = builder.AddInstruction( - HloInstruction::CreateParameter(2, second_weight_shape, "second_weight")); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto first_dot = builder.AddInstruction( - HloInstruction::CreateDot(intermediate_shape, lhs, first_weight, - dot_dnums, DefaultPrecisionConfig(2))); - - auto second_dot = builder.AddInstruction( - HloInstruction::CreateDot(result_shape, first_dot, second_weight, - dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence( - computation, {lhs, first_weight, second_weight, first_dot, second_dot}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.max_cross_program_prefetches = -1; - options.max_size_in_bytes = 256; - options.alignment_in_bytes = 8; - AssignMemorySpace(module.get(), std::move(options)); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 2); - EXPECT_EQ(cross_program_prefetches[0].parameter, 1); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - EXPECT_EQ(cross_program_prefetches[1].parameter, 2); - EXPECT_EQ(cross_program_prefetches[1].index, ShapeIndex({})); - - EXPECT_THAT( - module->entry_computation()->root_instruction(), - op::Dot(op::Dot(op::Parameter(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::Parameter(1))), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::Parameter(2)))); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchTupleTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - auto tuple_shape = ShapeUtil::MakeTupleShape({lhs_shape, rhs_shape}); - HloInstruction* param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - - auto lhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(lhs_shape, param, 0)); - auto rhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(rhs_shape, param, 1)); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {param, lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({1})); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchBitcastTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kOutput, kFeature}); - auto bitcast_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - HloInstruction* lhs = builder.AddInstruction( - HloInstruction::CreateParameter(0, lhs_shape, "lhs")); - HloInstruction* rhs = builder.AddInstruction( - HloInstruction::CreateParameter(1, rhs_shape, "rhs")); - - auto bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(bitcast_shape, rhs)); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, bitcast, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {lhs, rhs, bitcast, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 1); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchBitcastTupleTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kOutput, kFeature}); - auto bitcast_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - auto tuple_shape = ShapeUtil::MakeTupleShape({lhs_shape, rhs_shape}); - HloInstruction* param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - - auto lhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(lhs_shape, param, 0)); - auto rhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(rhs_shape, param, 1)); - - auto bitcast = - builder.AddInstruction(HloInstruction::CreateBitcast(bitcast_shape, rhs)); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, bitcast, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {param, lhs, rhs, bitcast, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({1})); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchNestedTupleTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - auto tuple_shape = ShapeUtil::MakeTupleShape({lhs_shape, rhs_shape}); - auto tuple_tuple_shape = ShapeUtil::MakeTupleShape({tuple_shape}); - HloInstruction* param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_tuple_shape, "p0")); - - auto gte = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(tuple_shape, param, 0)); - - auto lhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(lhs_shape, gte, 0)); - auto rhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(rhs_shape, gte, 1)); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {param, gte, lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchUnusedParamTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - HloInstruction* param = builder.AddInstruction( - HloInstruction::CreateParameter(0, rhs_shape, "p0")); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {param}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchTooBigTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 8; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - HloInstruction* lhs = builder.AddInstruction( - HloInstruction::CreateParameter(0, lhs_shape, "lhs")); - HloInstruction* rhs = builder.AddInstruction( - HloInstruction::CreateParameter(1, rhs_shape, "rhs")); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchTooBigTupleTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 8; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - auto tuple_shape = ShapeUtil::MakeTupleShape({lhs_shape, rhs_shape}); - HloInstruction* param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - - auto lhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(lhs_shape, param, 0)); - auto rhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(rhs_shape, param, 1)); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {param, lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchFusionTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 2; - constexpr int kFeature = 2; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - - auto module = CreateNewVerifiedModule(); - HloComputation::Builder fusion_builder("fusion"); - { - HloInstruction* lhs = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(0, lhs_shape, "lhs")); - HloInstruction* rhs = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(1, rhs_shape, "rhs")); - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = fusion_builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - (void)dot; - } - HloComputation* fusion_computation = - module->AddEmbeddedComputation(fusion_builder.Build()); - - auto activations = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{0.0, 1.0}, {2.0, 3.0}}))); - auto weights = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{0.0, 1.0}, {2.0, 3.0}}))); - HloInstruction* fusion = builder.AddInstruction(HloInstruction::CreateFusion( - result_shape, HloInstruction::FusionKind::kCustom, {activations, weights}, - fusion_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {activations, weights, fusion}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchFusionTupleTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 2; - constexpr int kFeature = 2; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShape(F32, {kFeature, kOutput}); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - auto tuple_shape = ShapeUtil::MakeTupleShape({lhs_shape, rhs_shape}); - - auto module = CreateNewVerifiedModule(); - HloComputation::Builder fusion_builder("fusion"); - { - HloInstruction* param = fusion_builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - auto lhs = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(lhs_shape, param, 0)); - auto rhs = fusion_builder.AddInstruction( - HloInstruction::CreateGetTupleElement(rhs_shape, param, 1)); - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = fusion_builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - (void)dot; - } - HloComputation* fusion_computation = - module->AddEmbeddedComputation(fusion_builder.Build()); - - auto activations = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{0.0, 1.0}, {2.0, 3.0}}))); - auto weights = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{0.0, 1.0}, {2.0, 3.0}}))); - HloInstruction* tuple = builder.AddInstruction( - HloInstruction::CreateTuple({activations, weights})); - HloInstruction* fusion = builder.AddInstruction(HloInstruction::CreateFusion( - result_shape, HloInstruction::FusionKind::kCustom, {tuple}, - fusion_computation)); - - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {activations, weights, tuple, fusion}); - CHECK_OK(module->set_schedule(schedule)); - - AssignMemorySpace(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchPinnedTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShapeWithDenseLayout( - F32, {kFeature, kOutput}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - HloInstruction* lhs = builder.AddInstruction( - HloInstruction::CreateParameter(0, lhs_shape, "lhs")); - HloInstruction* rhs = builder.AddInstruction( - HloInstruction::CreateParameter(1, rhs_shape, "rhs")); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.is_allowed_in_alternate_mem_fn = [](const HloValue& value) { - return true; - }; - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options)); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_GT(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchPinnedTupleTest) { - HloComputation::Builder builder(TestName()); - - constexpr int kBatch = 8; - constexpr int kFeature = 8; - constexpr int kOutput = 2; - - auto lhs_shape = ShapeUtil::MakeShape(F32, {kBatch, kFeature}); - auto rhs_shape = ShapeUtil::MakeShapeWithDenseLayout( - F32, {kFeature, kOutput}, - /*minor_to_major=*/{1, 0}, /*tiles=*/{}, - /*tail_padding_alignment_in_elements=*/1, /*element_size_in_bits=*/0, - kAlternateMemorySpace); - auto result_shape = ShapeUtil::MakeShape(F32, {kBatch, kOutput}); - auto tuple_shape = ShapeUtil::MakeTupleShape({lhs_shape, rhs_shape}); - HloInstruction* param = builder.AddInstruction( - HloInstruction::CreateParameter(0, tuple_shape, "p0")); - - auto lhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(lhs_shape, param, 0)); - auto rhs = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(rhs_shape, param, 1)); - - DotDimensionNumbers dot_dnums; - dot_dnums.add_lhs_contracting_dimensions(1); - dot_dnums.add_rhs_contracting_dimensions(0); - auto dot = builder.AddInstruction(HloInstruction::CreateDot( - result_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {param, lhs, rhs, dot}); - CHECK_OK(module->set_schedule(schedule)); - - Options options = DefaultMemorySpaceOptions(); - options.is_allowed_in_alternate_mem_fn = [](const HloValue& value) { - return true; - }; - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options)); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_GT(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootDupMayAlias) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true, input_output_alias={ {}: (0, {}, may-alias) } - ENTRY CrossProgramPrefetch { - c0 = s32[1,2] constant({{77, 77}}) - c1 = s32[] constant(0) - p0 = s32[2,2] parameter(0) - ROOT dup = s32[2,2] dynamic-update-slice(s32[2,2] p0, s32[1,2] c0, s32[] c1, s32[] c1) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); - EXPECT_THAT(FindInstruction(module.get(), "dup")->operand(0), - op::Parameter(0)); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootDusFusionMayAlias) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true, input_output_alias={ {}: (0, {}, may-alias) } - fused_computation { - fused_p0 = s32[2,2] parameter(0) - fused_p1 = s32[1,2] parameter(1) - fused_p2 = s32[] parameter(2) - fused_p3 = s32[] parameter(3) - ROOT dus = s32[2,2] dynamic-update-slice(fused_p0, fused_p1, fused_p2, fused_p3) - } - - ENTRY CrossProgramPrefetch { - p0 = s32[2,2] parameter(0) - c0 = s32[1,2] constant({{77, 77}}) - c1 = s32[] constant(0) - bitcast1 = s32[2,2] bitcast(p0) - ROOT fusion = s32[2,2] fusion(bitcast1, c0, c1, c1), kind=kLoop, calls=fused_computation - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootDup) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - ENTRY CrossProgramPrefetch { - c0 = s32[1,2] constant({{77, 77}}) - c1 = s32[] constant(0) - p0 = s32[2,2] parameter(0) - ROOT dup = s32[2,2] dynamic-update-slice(s32[2,2] p0, s32[1,2] c0, s32[] c1, s32[] c1) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); - EXPECT_THAT(FindInstruction(module.get(), "dup")->operand(0), - op::Parameter(0)); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootDupDot) { - // Cross program prefetch since the parameter and the root don't alias. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - ENTRY CrossProgramPrefetch { - c0 = s32[1,2] constant({{77, 77}}) - c1 = s32[] constant(0) - p0 = s32[2,2] parameter(0) - p1 = s32[2,2] parameter(1) - dup = s32[2,2] dynamic-update-slice(s32[2,2] p0, s32[1,2] c0, s32[] c1, s32[] c1) - ROOT dot = s32[2,2] dot(p1, dup), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_THAT(FindInstruction(module.get(), "dup")->operand(0), - op::AsyncCopy(kAlternateMemorySpace, kDefaultMemorySpace, - op::Parameter(0))); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootDotMayAlias) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true, input_output_alias={ {}: (0, {}, may-alias) } - ENTRY CrossProgramPrefetch { - p0 = s32[2,2] parameter(0) - p1 = s32[2,2] parameter(1) - ROOT dot = s32[2,2] dot(p1, p0), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); - EXPECT_THAT(FindInstruction(module.get(), "dot")->operand(1), - op::Parameter(0)); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootLiveOutBug) { - // Input-output aliased buffers should not be cross-program prefetched since - // the update on the buffer will not be reflected on the next program - // execution (the data in the alternate memory would be stale). - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias) } - fused_computation { - p0 = s32[2,2] parameter(0) - p1 = s32[2,2] parameter(1) - slice = s32[1,2] slice(p1), slice={[0:1], [0:2]} - c1 = s32[] constant(0) - ROOT dus = s32[2,2] dynamic-update-slice(s32[2,2] p0, s32[1,2] slice, s32[] c1, s32[] c1) - } - - ENTRY CrossProgramPrefetch { - p0 = s32[2,2] parameter(0) - p1 = s32[2,2] parameter(1) - dot = s32[2,2] dot(p1, p0), lhs_contracting_dims={0}, rhs_contracting_dims={0} - fusion = s32[2,2] fusion(p0, dot), kind=kLoop, calls=fused_computation - ROOT root = (s32[2,2], s32[2,2]) tuple(fusion, dot) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramRootParameter) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - ENTRY CrossProgramPrefetch { - p0 = s32[2,2] parameter(0) - ROOT bitcast = u32[2,2] bitcast(p0) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchNoReuse) { - // This test is for checking if the cross-program-prefetched buffer is freed - // after its last use and there is an end-of-program prefetch. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY CrossProgramPrefetch { - p0 = f32[8,8]{1,0} parameter(0) - p1 = f32[8,2]{1,0} parameter(1) - dot = f32[8,2]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT negate.9 = f32[8,2]{1,0} negate(negate.8) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto options = DefaultMemorySpaceOptions(); - // Enough space to fit the cross-program prefetch for both p0 and p1. - options.max_size_in_bytes = 512; - auto preset_assignments = AssignMemorySpace(module.get(), std::move(options), - /*max_prefetch_interval=*/5, - /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 1); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - LOG(INFO) << "module: " << module->ToString(); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(1), {}); - // Expect that there are two prefetches that use this value, one is the - // cross-program prefetch, the other is the end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 1); - // Also verify that the copy-done for the end-of-program prefetch is the last - // instruction in schedule. - const HloInstruction* last_instruction = - module->schedule() - .sequence(module->entry_computation()) - .instructions()[module->entry_computation()->instruction_count() - 1]; - EXPECT_THAT(last_instruction, op::CopyDone()); - EXPECT_NE(last_instruction, module->entry_computation()->root_instruction()); - // Cross program prefetch would use offset 0 because that's the first - // assignment. Since we are freeing the cross-program prefetch buffer, we - // would also expect to see some of the intermediate computations (one of the - // negate ops) to also get 0 offset allocations. - bool has_zero_offset_allocations = false; - for (auto pos_and_chunk : preset_assignments->chunks()) { - if (pos_and_chunk.first.instruction->opcode() == HloOpcode::kNegate && - pos_and_chunk.second.offset == 0) { - has_zero_offset_allocations = true; - } - } - EXPECT_TRUE(has_zero_offset_allocations); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchWithOverrideNoReuse) { - // This test is same as above, but with an override to cross-program prefetch - // parameter0 as opposed to p0 and limiting the max alternate memory - // size to 256 bytes so that both p0 and p1 cannot be assigned to alternate - // memory and priority is given to p0. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY CrossProgramPrefetch { - p0 = f32[8,8]{1,0} parameter(0) - p1 = f32[8,2]{1,0} parameter(1) - dot = f32[8,2]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT negate.9 = f32[8,2]{1,0} negate(negate.8) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto options = DefaultMemorySpaceOptions(); - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { - instruction_name_regex: "p(.*)" - instruction_regex: ".*parameter\\(0\\).*" - } - override_options { assign_first: true } - apply_to_cross_program_prefetches: true - })pb"; - ASSERT_OK_AND_ASSIGN(options.msa_sort_order_overrides, - ParseTextProto(text_proto)); - options.max_size_in_bytes = 256; - auto preset_assignments = AssignMemorySpace(module.get(), std::move(options), - /*max_prefetch_interval=*/5, - /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - LOG(INFO) << "module: " << module->ToString(); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(0), {}); - // Expect that there are two prefetches that use this value, one is the - // cross-program prefetch, the other is the end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 1); - // Also verify that the copy-done for the end-of-program prefetch is the last - // instruction in schedule. - const HloInstruction* last_instruction = - module->schedule() - .sequence(module->entry_computation()) - .instructions()[module->entry_computation()->instruction_count() - 1]; - EXPECT_THAT(last_instruction, op::CopyDone()); - EXPECT_NE(last_instruction, module->entry_computation()->root_instruction()); - // Cross program prefetch would use offset 0 because that's the first - // assignment. Since we are freeing the cross-program prefetch buffer, we - // would also expect to see some of the intermediate computations (one of the - // negate ops) to also get 0 offset allocations. - bool has_zero_offset_allocations = false; - for (auto pos_and_chunk : preset_assignments->chunks()) { - if (pos_and_chunk.first.instruction->opcode() == HloOpcode::kNegate && - pos_and_chunk.second.offset == 0) { - has_zero_offset_allocations = true; - } - } - EXPECT_TRUE(has_zero_offset_allocations); -} - -TEST_F(MemorySpaceAssignmentTest, UserAnnotatedCrossProgramPrefetchNoReuse) { - // This test is same as above, but with user directive to cross-program - // prefetch parameter0 as opposed to p0 and limiting the max alternate memory - // size to 256 bytes so that both p0 and p1 cannot be assigned to alternate - // memory and priority is given to p0. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true, entry_computation_layout={(f32[8,8]{1,0:S(1)}, f32[8,2]{1,0})->f32[8,2]{1,0}} - - ENTRY CrossProgramPrefetch { - p0 = f32[8,8]{1,0:S(1)} parameter(0) - p1 = f32[8,2]{1,0} parameter(1) - dot = f32[8,2]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT negate.9 = f32[8,2]{1,0} negate(negate.8) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 256; - auto preset_assignments = AssignMemorySpace(module.get(), std::move(options), - /*max_prefetch_interval=*/5, - /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - LOG(INFO) << "module: " << module->ToString(); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(0), {}); - // Expect that there are two prefetches that use this value, one is the - // cross-program prefetch, the other is the end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 1); - // Also verify that the copy-done for the end-of-program prefetch is the last - // instruction in schedule. - const HloInstruction* last_instruction = - module->schedule() - .sequence(module->entry_computation()) - .instructions()[module->entry_computation()->instruction_count() - 1]; - EXPECT_THAT(last_instruction, op::CopyDone()); - EXPECT_NE(last_instruction, module->entry_computation()->root_instruction()); - // Cross program prefetch would use offset 0 because that's the first - // assignment. Since we are freeing the cross-program prefetch buffer, we - // would also expect to see some of the intermediate computations (one of the - // negate ops) to also get 0 offset allocations. - bool has_zero_offset_allocations = false; - for (auto pos_and_chunk : preset_assignments->chunks()) { - if (pos_and_chunk.first.instruction->opcode() == HloOpcode::kNegate && - pos_and_chunk.second.offset == 0) { - has_zero_offset_allocations = true; - } - } - EXPECT_TRUE(has_zero_offset_allocations); - XLA_VLOG_LINES(3, module->ToString()); - bool found = false; - for (auto* c : module->computations()) { - for (auto* instr : c->instructions()) { - if (instr->name() == "p0") { - found = true; - EXPECT_EQ(instr->shape().layout().memory_space(), 0); - EXPECT_EQ(module->entry_computation_layout() - .parameter_layout(0) - .shape() - .layout() - .memory_space(), - 0); - } - } - } - EXPECT_TRUE(found); -} - -TEST_F(MemorySpaceAssignmentTest, - UserAnnotatedCrossProgramPrefetchWithoutPropagationToParameterNoReuse) { - // This test is same as above, but the S(1) memory space specified in the - // layout to cross-program prefetch p0 is only present in the entry - // computation layout and has not been propagated to the parameter - // instruction. This still works as the previous test. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true, entry_computation_layout={(f32[8,8]{1,0:S(1)}, f32[8,2]{1,0})->f32[8,2]{1,0}} - - ENTRY CrossProgramPrefetch { - p0 = f32[8,8]{1,0} parameter(0) - p1 = f32[8,2]{1,0} parameter(1) - dot = f32[8,2]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT negate.9 = f32[8,2]{1,0} negate(negate.8) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 256; - auto preset_assignments = AssignMemorySpace(module.get(), std::move(options), - /*max_prefetch_interval=*/5, - /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - LOG(INFO) << "module: " << module->ToString(); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(0), {}); - // Expect that there are two prefetches that use this value, one is the - // cross-program prefetch, the other is the end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 1); - // Also verify that the copy-done for the end-of-program prefetch is the last - // instruction in schedule. - const HloInstruction* last_instruction = - module->schedule() - .sequence(module->entry_computation()) - .instructions()[module->entry_computation()->instruction_count() - 1]; - EXPECT_THAT(last_instruction, op::CopyDone()); - EXPECT_NE(last_instruction, module->entry_computation()->root_instruction()); - // Cross program prefetch would use offset 0 because that's the first - // assignment. Since we are freeing the cross-program prefetch buffer, we - // would also expect to see some of the intermediate computations (one of the - // negate ops) to also get 0 offset allocations. - bool has_zero_offset_allocations = false; - for (auto pos_and_chunk : preset_assignments->chunks()) { - if (pos_and_chunk.first.instruction->opcode() == HloOpcode::kNegate && - pos_and_chunk.second.offset == 0) { - has_zero_offset_allocations = true; - } - } - EXPECT_TRUE(has_zero_offset_allocations); - XLA_VLOG_LINES(3, module->ToString()); - bool found = false; - for (auto* c : module->computations()) { - for (auto* instr : c->instructions()) { - if (instr->name() == "p0") { - found = true; - EXPECT_EQ(instr->shape().layout().memory_space(), 0); - EXPECT_EQ(module->entry_computation_layout() - .parameter_layout(0) - .shape() - .layout() - .memory_space(), - 0); - } - } - } - EXPECT_TRUE(found); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchTupleNoReuse) { - // This test is for checking if the cross-program-prefetched buffer is freed - // after its last use and there is an end-of-program prefetch. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY CrossProgramPrefetch { - p0 = (f32[8,8]{1,0}, f32[8,2]{1,0}) parameter(0) - get-tuple-element = f32[8,8]{1,0} get-tuple-element(p0), index=0 - get-tuple-element.1 = f32[8,2]{1,0} get-tuple-element(p0), index=1 - dot = f32[8,2]{1,0} dot(get-tuple-element, get-tuple-element.1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT negate.9 = f32[8,2]{1,0} negate(negate.8) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpace( - module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({1})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(0), {1}); - // Expect that there are two prefetches that use this value, one is the - // cross-program prefetch, the other is the end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 1); - // Also verify that the copy-done for the end-of-program prefetch is the last - // instruction in schedule. - const HloInstruction* last_instruction = - module->schedule() - .sequence(module->entry_computation()) - .instructions()[module->entry_computation()->instruction_count() - 1]; - EXPECT_THAT(last_instruction, op::CopyDone()); - EXPECT_NE(last_instruction, module->entry_computation()->root_instruction()); - // Cross program prefetch would use offset 0 because that's the first - // assignment. Since we are freeing the cross-program prefetch buffer, we - // would also expect to see some of the intermediate computations (one of the - // negate ops) to also get 0 offset allocations. - bool has_zero_offset_allocations = false; - for (auto pos_and_chunk : preset_assignments->chunks()) { - if (pos_and_chunk.first.instruction->opcode() == HloOpcode::kNegate && - pos_and_chunk.second.offset == 0) { - has_zero_offset_allocations = true; - } - } - EXPECT_TRUE(has_zero_offset_allocations); -} - -TEST_F(MemorySpaceAssignmentTest, - CrossProgramPrefetchEndOfProgramPrefetchAndWhile) { - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - while_condition { - param1 = (f32[8,2]{1,0}, f32[8,2]{1,0}) parameter(0) - ROOT cond = pred[] constant(true) - } - - while_body { - param2 = (f32[8,2]{1,0}, f32[8,2]{1,0}) parameter(0) - gte2 = f32[8,2]{1,0} get-tuple-element(param2), index=0 - gte3 = f32[8,2]{1,0} get-tuple-element(param2), index=1 - add = f32[8,2]{1,0} add(gte2, gte3) - negate.2 = f32[8,2]{1,0} negate(add) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT tuple2 = (f32[8,2]{1,0}, f32[8,2]{1,0}) tuple(negate.8, gte3) - } - - ENTRY CrossProgramPrefetch { - p0 = f32[8,8]{1,0} parameter(0) - p1 = f32[8,2]{1,0} parameter(1) - dot = f32[8,2]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - tuple = (f32[8,2]{1,0}, f32[8,2]{1,0}) tuple(negate.1, dot) - while = (f32[8,2]{1,0}, f32[8,2]{1,0}) while(tuple), condition=while_condition, body=while_body - ROOT gte0 = f32[8,2]{1,0} get-tuple-element(while), index=0 - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - auto preset_assignments = AssignMemorySpaceUsingCostAnalysis(module.get()); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 1); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - LOG(INFO) << "module: " << module->ToString(); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(1), {}); - // Expect that there are two prefetches that use this value, one is the - // cross-program prefetch, the other is the end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 1); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchReuse) { - // This tests the scenario that the cross-program-prefetched buffer is used - // again close to the end of the computation. In this case, it is better not - // to free the buffer. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY CrossProgramPrefetch { - p0 = f32[8,8]{1,0} parameter(0) - p1 = f32[8,2]{1,0} parameter(1) - dot = f32[8,2]{1,0} dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT dot.2 = f32[2,2]{1,0} dot(negate.8, p1), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 1); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(1), {}); - // Expect that there is one prefetch that use this value, the cross-program - // prefetch. There shouldn't be an end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchTupleReuse) { - // This tests the scenario that the cross-program-prefetched buffer is used - // again close to the end of the computation. In this case, it is better not - // to free the buffer. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY CrossProgramPrefetch { - p0 = (f32[8,8]{1,0}, f32[8,2]{1,0}) parameter(0) - get-tuple-element = f32[8,8]{1,0} get-tuple-element(p0), index=0 - get-tuple-element.1 = f32[8,2]{1,0} get-tuple-element(p0), index=1 - dot = f32[8,2]{1,0} dot(get-tuple-element, get-tuple-element.1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot) - negate.2 = f32[8,2]{1,0} negate(negate.1) - negate.3 = f32[8,2]{1,0} negate(negate.2) - negate.4 = f32[8,2]{1,0} negate(negate.3) - negate.5 = f32[8,2]{1,0} negate(negate.4) - negate.6 = f32[8,2]{1,0} negate(negate.5) - negate.7 = f32[8,2]{1,0} negate(negate.6) - negate.8 = f32[8,2]{1,0} negate(negate.7) - ROOT dot.2 = f32[2,2]{1,0} dot(negate.8, get-tuple-element.1), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - AssignMemorySpace(module.get(), DefaultMemorySpaceOptions(), - /*max_prefetch_interval=*/5, /*min_prefetch_interval=*/2); - - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); - EXPECT_EQ(cross_program_prefetches[0].parameter, 0); - EXPECT_EQ(cross_program_prefetches[0].index, ShapeIndex({1})); - - ASSERT_OK_AND_ASSIGN(std::unique_ptr dataflow_analysis, - HloDataflowAnalysis::Run(*module)); - const HloValue& cross_program_prefetched_value = - dataflow_analysis->GetValueDefinedAt( - module->entry_computation()->parameter_instruction(0), {1}); - // Expect that there is one prefetch that use this value, the cross-program - // prefetch. There shouldn't be an end-of-program prefetch. - auto is_cross_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_cross_program_prefetch), - 1); - auto is_end_of_program_prefetch = [](const HloUse& use) { - return use.instruction->opcode() == HloOpcode::kCopyStart && - !use.instruction->cross_program_prefetch_index().has_value(); - }; - EXPECT_EQ(absl::c_count_if(cross_program_prefetched_value.GetUses(), - is_end_of_program_prefetch), - 0); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchBufferUnused) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -%fused_computation { - %param_0.2 = f32[32]{0} parameter(0) - %param_1.4 = s32[100]{0} parameter(1) - %custom-call.1 = s32[100]{0} custom-call(s32[100]{0} %param_1.4), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[100]{0}} - %slice.1 = s32[32]{0} slice(s32[100]{0} %custom-call.1), slice={[0:32]} - %reshape.7 = s32[32]{0} reshape(s32[32]{0} %slice.1) - %transpose.5 = s32[32]{0} transpose(s32[32]{0} %reshape.7), dimensions={0} - %gather.1 = f32[32]{0} gather(f32[32]{0} %param_0.2, s32[32]{0} %transpose.5), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1} - %transpose.4 = f32[32]{0} transpose(f32[32]{0} %gather.1), dimensions={0} - ROOT %reshape.6 = f32[32]{0} reshape(f32[32]{0} %transpose.4) -} - -%i.reduce_sub_computation { - %rhs = s32[] parameter(1) - %lhs = s32[] parameter(0) - ROOT %add = s32[] add(s32[] %lhs, s32[] %rhs) -} - -%fused_computation.1 { - %constant.4 = s32[] constant(0) - %broadcast.4 = s32[100]{0} broadcast(s32[] %constant.4), dimensions={} - %param_0.4 = s32[32]{0} parameter(0) - %pad.1 = s32[100]{0} pad(s32[32]{0} %param_0.4, s32[] %constant.4), padding=0_68 - %constant.3 = s32[] constant(76031) - %broadcast.3 = s32[100]{0} broadcast(s32[] %constant.3), dimensions={} - ROOT %clamp.1 = s32[100]{0} clamp(s32[100]{0} %broadcast.4, s32[100]{0} %pad.1, s32[100]{0} %broadcast.3) -} - -ENTRY %main { - %constant = s32[] constant(0) - %i = s32[32,1]{0,1} parameter(1) - %o = f32[32]{0} parameter(0) - %reduce = s32[32]{0} reduce(s32[32,1]{0,1} %i, s32[] %constant), dimensions={1}, to_apply=%i.reduce_sub_computation - %fusion.1 = s32[100]{0} fusion(s32[32]{0} %reduce), kind=kLoop, calls=%fused_computation.1 - ROOT %fusion = f32[32]{0} fusion(f32[32]{0} %o, s32[100]{0} %fusion.1), kind=kCustom, calls=%fused_computation -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Fusion(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, op::Parameter(0)), - op::Fusion())); -} - -TEST_F(MemorySpaceAssignmentTest, CrossProgramPrefetchPermissiveMode) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -fused_computation { - param_0 = f32[2] parameter(0) - param_1 = f32[4,2] parameter(1) - broadcast = f32[4,2] broadcast(param_0), dimensions={1} - ROOT multiply = f32[4,2] multiply(broadcast, param_1) -} - -ENTRY entry { - p0 = f32[2] parameter(0) - p1 = f32[4,2] parameter(1) - fusion = f32[4,2] fusion(p0, p1), kind=kLoop, calls=fused_computation - ROOT negate = f32[4,2] negate(fusion) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.cross_program_prefetch_permissive_mode = true; - AssignMemorySpace(module.get(), std::move(options)); - auto cross_program_prefetches = module->CrossProgramPrefetches(); - EXPECT_EQ(cross_program_prefetches.size(), 1); -} - -// Test description: -// - Setup: Make sure p1 can not be prefetched to alternate memory until after -// instruction c. We do this by causing p0 to be prefetched to alternate -// memory for use in c. Since p0 is larger than 1/2 of alternate memory, we -// will not be able to prefetch p1 until after p0 is unallocated. -// - Test: prefetch p1, after p0 is unallocated from alternate memory (after -// instruction c). -TEST_F(MemorySpaceAssignmentTest, CopyResourceIntegration) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY main { - p0 = s32[8,8] parameter(0) - p1 = s32[8,8] parameter(1) - p2 = s32[] parameter(2) - a = negate(p2) - b = negate(a) - c = add(p0, p0) - d = negate(b) - e = negate(d) - f = add(p1, p1) - - ROOT result = tuple(e,c,f) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 300; - - // Setup cost analysis so it takes 2 instructions to prefetch anything. - HloCostAnalysis::Properties properties; - properties[HloCostAnalysis::kBytesAccessedKey] = kBytesPerSecond; - HloCostAnalysis hlo_cost_analysis(HloCostAnalysis::DefaultShapeSize, - properties); - HloCostAnalysisWithAcceptState hlo_cost_analysis_wrapper(hlo_cost_analysis); - CostAnalysisOptions cost_analysis_options; - cost_analysis_options.default_mem_bandwidth_bytes_per_second = - kBytesPerSecond; - OpCostManager op_cost_manager( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - ASSERT_OK_AND_ASSIGN(auto cost_analysis, - FakeCostAnalysis::Create(op_cost_manager, *module, - cost_analysis_options)); - cost_analysis->SetOverrideForGetInstructionElapsed( - [](const HloInstruction& instruction) -> float { return 10.0; }); - cost_analysis->SetOverrideForGetAsyncCopyElapsed( - [](int64_t size_in_bytes) -> float { return 20.0; }); - options.cost_analysis = cost_analysis.get(); - CostAnalysisPrefetchIntervalPicker prefetch_interval_picker( - CostAnalysisPrefetchIntervalPicker( - *cost_analysis, /*min_overlap_to_async_copy_ratio=*/0.8, - /*preferred_overlap_to_async_copy_ratio=*/1.5, - /*max_overlap_to_mem_size_async_copy_ratio=*/10.0, - /*mem_size_bytes=*/options.max_size_in_bytes)); - - // p0 has the highest priority, followed by p1, followed by everything else. - MsaBufferIntervalCompare compare = [](const MsaBufferInterval& lhs, - const MsaBufferInterval& rhs) -> bool { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that for p0 and p1. - int priority = 100; - if (x.buffer->instruction()->name() == "p0") { - priority = 0; - } else if (x.buffer->instruction()->name() == "p1") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - - // Run test. - AssignMemorySpace(module.get(), std::move(options), compare, - &prefetch_interval_picker); - - // - Make sure the setup occurred, i.e., that p0 is prefetched to alternate - // memory for use by c. - // - Make sure p1 is prefetched. - ASSERT_THAT( - module->entry_computation()->root_instruction(), - op::Tuple(_, - // p0 is prefetched to alternate memory for use by c. - op::Add(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, op::Parameter(0)), - op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, op::Parameter(0))), - // p1 is prefetched to alternate memory for use by f. - op::Add(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, op::Parameter(1)), - op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, op::Parameter(1))))); - - // Check the schedule - const std::vector& schedule = - module->schedule().sequence(module->entry_computation()).instructions(); - auto find_schedule_index = [&schedule](absl::string_view name) -> int { - for (int i = 0; i < schedule.size(); ++i) { - if (schedule[i]->name() == name) { - return i; - } - } - LOG(FATAL) << "Unable to find index of instruction with name " << name; - }; - int c_index = find_schedule_index("c"); - int p1_copy_start = find_schedule_index(module->entry_computation() - ->root_instruction() // result - ->operand(2) // f - ->operand(0) // copy done - ->operand(0) // copy start - ->name()); - int d_index = find_schedule_index("d"); - int e_index = find_schedule_index("e"); - int p1_copy_end = find_schedule_index(module->entry_computation() - ->root_instruction() // result - ->operand(2) // f - ->operand(0) // copy done - ->name()); - int f_index = find_schedule_index("f"); - // We expect to start copying p1 after c. - EXPECT_EQ(p1_copy_start, c_index + 1); - // d and e should follow come between p1's copy start and end. - EXPECT_EQ(d_index, p1_copy_start + 1); - EXPECT_EQ(e_index, d_index + 1); - EXPECT_EQ(p1_copy_end, e_index + 1); - // f should immediately follow the end of p1's copy. - EXPECT_EQ(f_index, p1_copy_end + 1); -} - -TEST_F(MemorySpaceAssignmentTest, ExpandScopedAlternateMemory) { - absl::string_view hlo_string = R"hlo( - HloModule TestModule, is_scheduled=true - ENTRY Main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - p2 = f32[8,8] parameter(2) - p3 = f32[8,8] parameter(3) - - v0 = add(p0, p1) - v1 = add(v0, p1) - v2 = add(v1, p1) - - v3 = multiply(v2, p2) - v4 = multiply(v3, p3) - - ROOT t = tuple(v3, v4) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that used for 'prefetch'. - int priority = 100; - if (x.buffer->instruction()->name() == "p2") { - priority = 1; - } else if (x.buffer->instruction()->name() == "p3") { - priority = 2; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 1000); - - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 600; - options.reserved_scoped_memory_fn = - [](const HloInstruction* instruction, - const absl::flat_hash_set< - std::pair>& /*operands_in_alternate_memory*/, - const absl::flat_hash_set< - ShapeIndex>& /*outputs_in_alternate_memory*/) { return 10; }; - options.expanded_scoped_alternate_memory_mode = - ExpandedScopedAlternateMemoryMode::ENABLED; - options.alignment_in_bytes = 10; - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - - VLOG(1) << "Post-MSA module:\n" << module->ToString(); - - // We expect MSA to do the following: - // A. Initially allocate [0, 10) for scoped alternate memory, for each - // instruction. - // B. Since, p2 comes first in the buffer sorting, we expect it to be - // allocated [10, 266) for a prefetch - // C. Since, p3 comes next in the buffer sorting, we expect it to be allocated - // [270, 526) for a prefetch - // D. Finally, MSA will try to expand the scoped alternate memory allocations - // to the largest available buffers, keeping in mind the prefetches. - - // Check B and C. - for (const auto& [position, chunk] : preset_assignments->chunks()) { - if (position.instruction->opcode() == HloOpcode::kCopyDone) { - ASSERT_EQ(position.instruction->operand_count(), 1); - const HloInstruction* copy_start = position.instruction->operand(0); - ASSERT_EQ(copy_start->operand_count(), 1); - const HloInstruction* copy_operand = copy_start->operand(0); - if (copy_operand->name() == "p2") { - EXPECT_EQ(chunk.offset, 10); - EXPECT_EQ(chunk.size, 256); - } else if (copy_operand->name() == "p3") { - EXPECT_EQ(chunk.offset, 270); - EXPECT_EQ(chunk.size, 256); - } - } - } - - // Check D. - for (const auto& [instruction, chunk] : - preset_assignments->scoped_allocation_chunks()) { - if (instruction->name() == "p0") { - // Extended scoped allocation. - EXPECT_EQ(chunk.offset, 0); - EXPECT_EQ(chunk.size, 600); - } else if (instruction->name() == "p1") { - // Extended scoped allocation. - EXPECT_EQ(chunk.offset, 0); - EXPECT_EQ(chunk.size, 600); - } else if (instruction->name() == "p2") { - // Extended scoped allocation. - EXPECT_EQ(chunk.offset, 0); - EXPECT_EQ(chunk.size, 600); - } else if (instruction->name() == "p3") { - // Moved scoped allocation. - EXPECT_EQ(chunk.offset, 270); - EXPECT_EQ(chunk.size, 330); - } else if (instruction->name() == "v0") { - // Moved scoped allocation. - EXPECT_EQ(chunk.offset, 530); - EXPECT_EQ(chunk.size, 70); - } else if (instruction->name() == "v1") { - // Moved scoped allocation. - EXPECT_EQ(chunk.offset, 530); - EXPECT_EQ(chunk.size, 70); - } else if (instruction->name() == "v2") { - // Moved scoped allocation. - EXPECT_EQ(chunk.offset, 530); - EXPECT_EQ(chunk.size, 70); - } else if (instruction->name() == "v3") { - // Moved scoped allocation. - EXPECT_EQ(chunk.offset, 530); - EXPECT_EQ(chunk.size, 70); - } else if (instruction->name() == "v4") { - // Extended scoped allocation. - EXPECT_EQ(chunk.offset, 0); - EXPECT_EQ(chunk.size, 270); - } else if (instruction->name() == "t") { - // Extended scoped allocation. - EXPECT_EQ(chunk.offset, 0); - EXPECT_EQ(chunk.size, 600); - } - } -} - -class SlicedPrefetchTest : public MemorySpaceAssignmentTestBase { - protected: - // Used by CheckSchedule() to classify instructions in the schedule. - enum class InstructionClass { - kUnknown, - // A slice start that we care about, as determined by the test. - kRelatedSliceStart, - // A slice done that we care about, as determined by the test. - kRelatedSliceDone, - // A concat-bitcast that we care about, as determined by the test. - kRelatedConcatBitcast, - // A non-copy-like instruction that slices should start after. - kStartAfterNonCopy, - // A non-copy-like instruction that slices should finish before. - kDoneBeforeNonCopy, - // A copy, slice, or concat-bitcast that we don't care about. - kUnrelatedCopyLike, - // All other instructions. - kUnrelatedNonCopy, - }; - - static std::string InstructionClassToString( - InstructionClass instruction_class) { - switch (instruction_class) { - case InstructionClass::kUnknown: - return "unknown"; - case InstructionClass::kRelatedSliceStart: - return "slice start"; - case InstructionClass::kRelatedSliceDone: - return "slice done"; - case InstructionClass::kRelatedConcatBitcast: - return "concat-bitcast"; - case InstructionClass::kStartAfterNonCopy: - return "start after non-copy"; - case InstructionClass::kDoneBeforeNonCopy: - return "done before non-copy"; - case InstructionClass::kUnrelatedCopyLike: - return "unrelated copy-like"; - case InstructionClass::kUnrelatedNonCopy: - return "unrelated non-copy"; - } - } - - // A class that can be mocked to set expectations on slice proposals. To do - // that, we set Options::propose_slice_fn to a lambda that calls our mocks - // ProposeSlices() method. - class SliceProposer { - public: - SliceProposer() = default; - virtual ~SliceProposer() = default; - - virtual absl::StatusOr ProposeSlices( - const Shape& shape, const SlicedPrefetchOptions& options) = 0; - }; - - class MockSliceProposer : public SliceProposer { - public: - MOCK_METHOD(absl::StatusOr, ProposeSlices, - (const Shape& shape, const SlicedPrefetchOptions& options), - (override)); - }; - - // An HloInstruction* matcher for matching the asynchronous sliced copies - // produced by MSA. In particular, the matcher performs the following - // checks: - // - The copy is concluded with a concat-bitcast custom call, or a - // bitcast of a concat-bitcast custom call if expect_bitcasted_io is true - // - The operands to the concat-bitcast are asynchronous slices of the - // expected operand, or asynchronous slices of a bitcast of the expected - // operand if expect_bitcasted_io is true - // - The number of slices is as expected (i.e., - // expected_slice_params_per_slice_in_spatial_order_.size()) - // - The copy is from and to the correct memory spaces - // - The shape before and after the copy is the same - // - When the slices are sorted in expected spatial order, their slice - // starts and limits are as expected - // - The slices are to the correct memory space - // - All slices have slice strides of 1 - class AsyncSlicedCopy - : public ::testing::MatcherInterface { - public: - // The parameters in expected_slice_params_per_slice_in_spatial_order should - // be sorted in the order in which we expect the corresponding slices to be - // laid out in memory. - AsyncSlicedCopy(int64_t to_space, int64_t from_space, - std::vector> - expected_slice_params_per_slice_in_spatial_order, - ::testing::Matcher operand, - bool expect_bitcasted_io) - : to_space_(to_space), - from_space_(from_space), - expected_slice_params_per_slice_in_spatial_order_( - std::move(expected_slice_params_per_slice_in_spatial_order)), - base_hlo_matcher_(CreateBaseHloMatcher( - operand, expected_slice_params_per_slice_in_spatial_order_.size(), - expect_bitcasted_io)), - expect_bitcasted_io_(expect_bitcasted_io) {} - - bool MatchAndExplain( - const HloInstruction* instruction, - ::testing::MatchResultListener* listener) const override { - // Match opcodes and number of operands. - if (!base_hlo_matcher_.MatchAndExplain(instruction, listener)) { - return false; - } - - // Check if the copied result has the proper memory space. - if (!MatchMemorySpace(instruction, to_space_, "copy result", listener)) { - return false; - } - - // Find some instructions in the async copy. - const HloInstruction* concat_bitcast = - (expect_bitcasted_io_ ? instruction->operand(0) : instruction); - VLOG(2) << "AsyncSlicedCopy identified the concat-bitcast as " - << concat_bitcast->name(); - const HloInstruction* copy_operand = - concat_bitcast->operand(0)->operand(0)->operand(0); - const HloInstruction* original_copy_operand = - (expect_bitcasted_io_ ? copy_operand->operand(0) : copy_operand); - VLOG(2) << "AsyncSlicedCopy identified the copy operand as " - << copy_operand->name() << ", and the original copy operand as " - << original_copy_operand->name(); - - // Check if the copied tensor has the proper memory space. - if (!MatchMemorySpace(original_copy_operand, from_space_, "copy operand", - listener)) { - return false; - } - - // Check if the copied tensor retains its shape. - if (!Shape::Equal().IgnoreMemorySpaceInLayout()( - instruction->shape(), original_copy_operand->shape())) { - *listener << " has a shape of " - << original_copy_operand->shape().ToString( - /*print_layout=*/true) - << " before copying but a shape of " - << instruction->shape().ToString(/*print_layout=*/true) - << " after copying (ignoring memory space)"; - - return false; - } - - // This should already be checked in the custom call matcher. - CHECK_EQ(concat_bitcast->operand_count(), - expected_slice_params_per_slice_in_spatial_order_.size()); - - // Check if the slicing parameters are correct and if the slices are to - // the correct memory space. - std::vector sorted_slices = - SortSlicesInExpectedSpatialOrder(concat_bitcast); - for (int i = 0; i < sorted_slices.size(); ++i) { - const HloInstruction* slice = - sorted_slices[i]->async_wrapped_instruction(); - - if (!MatchMemorySpace(slice, to_space_, "slice", listener)) { - return false; - } - - const std::vector& expected_slice_params_per_dim = - expected_slice_params_per_slice_in_spatial_order_[i]; - if (slice->slice_starts().empty()) { - *listener << " has slice (" << slice->name() - << "), with no slicing parameters"; - return false; - } - if (slice->slice_limits().size() != slice->slice_starts().size() || - slice->slice_strides().size() != slice->slice_limits().size()) { - *listener - << " has slice (" << slice->name() - << "), with an inconsistent number slice starts/limits/strides"; - return false; - } - if (slice->slice_starts().size() != - copy_operand->shape().dimensions().size()) { - *listener - << " has slice (" << slice->name() << "), with " - << slice->slice_starts().size() - << " slice parameters (i.e., starts/limits/strides), expected " - << expected_slice_params_per_slice_in_spatial_order_.size(); - return false; - } - for (int dim = 0; dim < slice->slice_starts().size(); ++dim) { - const SliceParam& expected_slice_params = - expected_slice_params_per_dim[dim]; - if (slice->slice_starts()[dim] != - expected_slice_params.start_inclusive) { - *listener << " has slice (" << slice->name() - << "), with slice start of " << slice->slice_starts()[dim] - << " at dim " << dim << ", expected " - << expected_slice_params.start_inclusive; - return false; - } - if (slice->slice_limits()[dim] != - expected_slice_params.end_exclusive) { - *listener << " has slice (" << slice->name() - << "), with slice limit of " << slice->slice_limits()[dim] - << " at dim " << dim << ", expected " - << expected_slice_params.end_exclusive; - return false; - } - if (slice->slice_strides()[dim] != 1) { - *listener << " has slice (" << slice->name() - << "), slice stride of " << slice->slice_strides()[dim] - << " at dim " << dim << ", expected 1"; - return false; - } - } - } - - return true; - } - - void DescribeTo(std::ostream* os) const override { - base_hlo_matcher_.DescribeTo(os); - std::vector slice_parameters_per_operand; - for (int op_idx = 0; - op_idx < expected_slice_params_per_slice_in_spatial_order_.size(); - ++op_idx) { - std::vector slice_params_per_dim; - for (int dim = 0; - dim < - expected_slice_params_per_slice_in_spatial_order_[op_idx].size(); - ++dim) { - const SliceParam& slice_params = - expected_slice_params_per_slice_in_spatial_order_[op_idx][dim]; - slice_params_per_dim.push_back(absl::StrCat( - "dim ", dim, ": {start: ", slice_params.start_inclusive, - ", limit: ", slice_params.end_exclusive, "}")); - } - slice_parameters_per_operand.push_back( - absl::StrCat("operand ", op_idx, ": { ", - absl::StrJoin(slice_params_per_dim, ", "), " }")); - } - *os << " (copying from memory space " << from_space_ << " to " - << to_space_ - << ", with asynchronous slice operands using the following slice " - "parameters: { " - << absl::StrJoin(slice_parameters_per_operand, ", ") << " })"; - } - - private: - static ::testing::Matcher CreateBaseHloMatcher( - ::testing::Matcher operand, int64_t num_slices, - bool expect_bitcasted_io) { - if (expect_bitcasted_io) { - return op::Bitcast(op::CustomCall( - kConcatBitcastCustomCall, - std::vector<::testing::Matcher>( - num_slices, - op::AsyncDone(op::AsyncStart(op::Bitcast(operand)))))); - } - return op::CustomCall( - kConcatBitcastCustomCall, - std::vector<::testing::Matcher>( - num_slices, op::AsyncDone(op::AsyncStart(operand)))); - } - - static bool MatchMemorySpace(const HloInstruction* instruction, - int64_t expected_memory_space, - absl::string_view error_message_identifier, - ::testing::MatchResultListener* listener) { - if (!instruction->shape().has_layout()) { - *listener << " contains " << error_message_identifier << " named " - << instruction->name() - << " without a layout, expected a layout with memory space " - << expected_memory_space; - return false; - } - if (instruction->shape().layout().memory_space() != - expected_memory_space) { - *listener << " contains " << error_message_identifier << " named " - << instruction->name() << " in memory space " - << expected_memory_space << ", expected " - << expected_memory_space; - return false; - } - - return true; - } - - int64_t to_space_; - int64_t from_space_; - std::vector> - expected_slice_params_per_slice_in_spatial_order_; - ::testing::Matcher base_hlo_matcher_; - bool expect_bitcasted_io_; - }; - - // Returns an AsyncSlicedCopy matcher. - static inline ::testing::Matcher IsAsyncSlicedCopy( - int64_t to_space, int64_t from_space, - std::vector> - expected_slice_params_per_slice_in_spatial_order, - ::testing::Matcher operand_matcher, - bool expect_bitcasted_io = false) { - return ::testing::MakeMatcher(new AsyncSlicedCopy( - to_space, from_space, expected_slice_params_per_slice_in_spatial_order, - operand_matcher, expect_bitcasted_io)); - } - - // We make our own matcher for SlicedPrefetchOptions to work around the fact - // third_party/tensorflow does not have any generic way to match proto - // buffers. - class SlicedPrefetchOptionsMatcher - : public ::testing::MatcherInterface { - public: - explicit SlicedPrefetchOptionsMatcher( - SlicedPrefetchOptions expected_options) - : expected_options_(std::move(expected_options)) {} - - bool MatchAndExplain( - const SlicedPrefetchOptions& options, - ::testing::MatchResultListener* listener) const override { - if (options.max_slices() != expected_options_.max_slices()) { - *listener << " has " << options.max_slices() << " max slices, expected " - << expected_options_.max_slices(); - return false; - } - - if (options.min_bytes() != expected_options_.min_bytes()) { - *listener << " has " << options.min_bytes() << " min bytes, expected " - << expected_options_.min_bytes(); - return false; - } - - if (options.fail_on_non_alignment_boundary_slice_proposal() != - expected_options_.fail_on_non_alignment_boundary_slice_proposal()) { - *listener - << " has fail_on_non_alignment_boundary_slice_proposal set to " - << options.fail_on_non_alignment_boundary_slice_proposal() - << ", expected " - << expected_options_ - .fail_on_non_alignment_boundary_slice_proposal(); - return false; - } - - return true; - } - - void DescribeTo(std::ostream* os) const override { - *os << " has the following options: max_slices(" - << expected_options_.max_slices() << "), min_bytes(" - << expected_options_.min_bytes() - << ") fail_on_non_alignment_boundary_slice_proposal(" - << expected_options_.fail_on_non_alignment_boundary_slice_proposal() - << ")"; - } - - private: - SlicedPrefetchOptions expected_options_; - }; - - // Returns an SlicedPrefetchOptions matcher. - static inline ::testing::Matcher - EqualsSlicedPrefetchOptions(SlicedPrefetchOptions expected_options) { - return ::testing::MakeMatcher( - new SlicedPrefetchOptionsMatcher(std::move(expected_options))); - } - - // Slices can be passed to the concat-bitcast in any order. This function - // sorts a the slices in the order they should spatially (in memory). Note, - // this function is specific to the way we are constructing slices for the - // test. E.g., it relies on the first dimension of the tensor to be the - // slice dimension. - // - // REQUIRES: - // - All operands of concat_bitcast must be asynchronous slices. - static std::vector SortSlicesInExpectedSpatialOrder( - const HloInstruction* concat_bitcast) { - std::vector sorted_slices( - concat_bitcast->operands().begin(), concat_bitcast->operands().end()); - - absl::c_sort(sorted_slices, [](const HloInstruction* lhs, - const HloInstruction* rhs) { - CHECK(IsAsyncSliceDone(lhs)); - CHECK(IsAsyncSliceDone(rhs)); - CHECK(!lhs->async_wrapped_instruction()->slice_starts().empty()); - CHECK(!rhs->async_wrapped_instruction()->slice_starts().empty()); - return lhs->async_wrapped_instruction()->slice_starts().front() < - rhs->async_wrapped_instruction()->slice_starts().front(); - }); - - return sorted_slices; - } - - // Returns true if instruction is an async copy start. - static bool IsAsyncCopyStart(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kCopyStart; - } - - // Returns true if instruction is an async copy done. - static bool IsAsyncCopyDone(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kCopyDone; - } - - // Returns true if instruction is an async slice start. - static bool IsAsyncSliceStart(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kAsyncStart && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice; - } - - // Returns true if instruction is an async slice done. - static bool IsAsyncSliceDone(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kAsyncDone && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice; - } - - // Returns true if instruction is a concat-bitcast. - static bool IsConcatBitcast(const HloInstruction* instruction) { - return instruction->IsCustomCall(kConcatBitcastCustomCall); - } - - // Returns the index of the first instruction with the given name. - static absl::StatusOr FindScheduleIndexOfInstruction( - const std::vector& schedule, absl::string_view name, - InstructionClass c) { - for (int i = 0; i < schedule.size(); ++i) { - if (schedule[i]->name() == name) { - return i; - } - } - - return NotFound( - "%s", - absl::StrCat("Could not find ", InstructionClassToString(c), - " instruction ", name, " in the instruction schedule.")); - } - - // Returns a scheduled instruction with the specified name or null. - static const HloInstruction* FindNamedScheduledInstruction( - const HloModule& module, absl::string_view name) { - for (const HloInstruction* i : module.entry_computation()->instructions()) { - if (i->name() == name) { - return i; - } - } - return nullptr; - } - - static absl::StatusOr> GetSliceStartIndices( - const std::vector& schedule, - const HloInstruction* concat_bitcast) { - std::vector indices; - - if (!IsConcatBitcast(concat_bitcast)) { - return InvalidArgumentStrCat(concat_bitcast->name(), - " is not a concat-bitcast."); - } - for (int i = 0; i < concat_bitcast->operand_count(); ++i) { - const HloInstruction* async_slice_done = concat_bitcast->operand(i); - if (!IsAsyncSliceDone(async_slice_done)) { - return InvalidArgumentStrCat("Operand ", i, " of ", - concat_bitcast->name(), - " is not an async-slice-done."); - } - const HloInstruction* async_slice_start = async_slice_done->operand(0); - if (!IsAsyncSliceStart(async_slice_start)) { - return InvalidArgumentStrCat("Operand 0, of operand ", i, " of ", - concat_bitcast->name(), - " is not an async-slice-start."); - } - ABSL_ASSIGN_OR_RETURN( - int schedule_index, - FindScheduleIndexOfInstruction(schedule, async_slice_start->name(), - InstructionClass::kRelatedSliceStart)); - indices.push_back(schedule_index); - } - - return indices; - } - - // REQUIRES: - // - Concat-bitcast and all slices were found in the schedule used to - // construct schedule_to_class. - static absl::Status ConcatBitcastAndSlicesAfterInstruction( - const std::vector& schedule, - const std::vector& schedule_to_class, - int slices_start_after_index) { - for (int i = 0; i < slices_start_after_index; ++i) { - InstructionClass c = schedule_to_class[i]; - const HloInstruction* instruction = schedule[i]; - - if (c == InstructionClass::kRelatedSliceStart || - c == InstructionClass::kRelatedSliceDone || - c == InstructionClass::kRelatedConcatBitcast) { - return FailedPrecondition( - "%s", absl::StrCat(InstructionClassToString(c), " ", - instruction->name(), " is scheduled at ", i, - ", but is expected to be after ", - schedule[slices_start_after_index]->name(), - " at ", slices_start_after_index, ".")); - } - } - - return absl::OkStatus(); - } - - // REQUIRES: - // - Concat-bitcast and all slices were found in the schedule used to - // construct schedule_to_class. - static absl::Status AtLeastOneNonCopyLikeInstructionBetweenSliceStarts( - const std::vector& schedule, - const std::vector& schedule_to_class) { - bool found_non_copy_since_last_slice_start = true; - for (int i = 0; i < schedule_to_class.size(); ++i) { - InstructionClass c = schedule_to_class[i]; - - if (c == InstructionClass::kRelatedSliceStart && - !found_non_copy_since_last_slice_start) { - return FailedPrecondition( - "%s", - absl::StrCat( - "Did not find a non-copy-like instruction between slice start ", - schedule[i]->name(), " at ", i, - " and the previous slice start.")); - } - - if (c == InstructionClass::kRelatedSliceStart) { - found_non_copy_since_last_slice_start = false; - } else if (c == InstructionClass::kUnrelatedNonCopy) { - found_non_copy_since_last_slice_start = true; - } - } - - return absl::OkStatus(); - } - - // REQUIRES: - // - Concat-bitcast and all slices were found in the schedule used to - // construct schedule_to_class. - static absl::Status OneSliceStartAfterInstructionWithNoCopyLikeBetween( - const std::vector& schedule, - const std::vector& schedule_to_class, - int slices_start_after_index) { - int first_slice_start_after_schedule_after = -1; - int first_non_copy_after_schedule_after = -1; - for (int i = slices_start_after_index + 1; - i < schedule_to_class.size() && - (first_slice_start_after_schedule_after == -1 || - first_non_copy_after_schedule_after == -1); - ++i) { - if (first_slice_start_after_schedule_after == -1 && - schedule_to_class[i] == InstructionClass::kRelatedSliceStart) { - first_slice_start_after_schedule_after = i; - continue; - } - if (first_non_copy_after_schedule_after == -1 && - schedule_to_class[i] == InstructionClass::kUnrelatedNonCopy) { - first_non_copy_after_schedule_after = i; - continue; - } - } - if (first_slice_start_after_schedule_after == -1) { - return NotFound( - "%s", absl::StrCat("Could not find a slice start instruction " - "after start after instruction ", - schedule[slices_start_after_index]->name(), " at ", - slices_start_after_index, ".")); - } - if (first_non_copy_after_schedule_after == -1) { - return NotFound( - "%s", absl::StrCat("Could not a find non-copy-like instruction " - "after start after instruction ", - schedule[slices_start_after_index]->name(), " at ", - slices_start_after_index, ".")); - } - if (first_slice_start_after_schedule_after > - first_non_copy_after_schedule_after) { - return FailedPrecondition( - "%s", absl::StrCat( - "Unexpectedly found a non-copy-like instruction at ", - first_non_copy_after_schedule_after, ", between ", - schedule[slices_start_after_index]->name(), " at ", - slices_start_after_index, ", and the first slice start at ", - first_slice_start_after_schedule_after, ".")); - } - - return absl::OkStatus(); - } - - // REQUIRES: - // - Concat-bitcast and all slices were found in the schedule used to - // construct schedule_to_class. - static absl::Status ConcatBitcastAndSlicesBeforeInstruction( - const std::vector& schedule, - const std::vector& schedule_to_class, - int slices_done_before_index) { - for (int i = slices_done_before_index + 1; i < schedule_to_class.size(); - ++i) { - InstructionClass c = schedule_to_class[i]; - const HloInstruction* instruction = schedule[i]; - - if (c == InstructionClass::kRelatedSliceStart || - c == InstructionClass::kRelatedSliceDone || - c == InstructionClass::kRelatedConcatBitcast) { - return FailedPrecondition( - "%s", absl::StrCat(InstructionClassToString(c), " ", - instruction->name(), " is scheduled at ", i, - ", but is expected to be before ", - schedule[slices_done_before_index]->name(), - " at ", slices_done_before_index, ".")); - } - } - - return absl::OkStatus(); - } - - // REQUIRES: - // - Concat-bitcast and all slices were found in the schedule used to - // construct schedule_to_class. - static absl::Status - ConcatBitcastAndSliceDonesBeforeInstructionWithNoCopyLikeBetween( - const std::vector& schedule, - const std::vector& schedule_to_class, - int slices_done_before_index) { - bool found_non_copy = false; - for (int i = slices_done_before_index - 1; i >= 0; --i) { - InstructionClass c = schedule_to_class[i]; - const HloInstruction* instruction = schedule[i]; - - if (c == InstructionClass::kUnrelatedNonCopy) { - found_non_copy = true; - continue; - } - - if (found_non_copy && (c == InstructionClass::kRelatedSliceDone || - c == InstructionClass::kRelatedConcatBitcast)) { - return FailedPrecondition( - "%s", - absl::StrCat("Found non-copy instruction between ", - InstructionClassToString(c), " ", instruction->name(), - " at ", i, ", and slice done before instruction ", - schedule[slices_done_before_index]->name(), " at ", - slices_done_before_index, ".")); - } - } - - return absl::OkStatus(); - } - - // REQUIRES: - // - Concat-bitcast and all slices were found in the schedule used to - // construct schedule_to_class. - static absl::Status ConcatBitcastAfterSliceDones( - const std::vector& schedule, - const std::vector& schedule_to_class) { - int concat_bitcast_index = -1; - for (int i = 0; i < schedule_to_class.size(); ++i) { - InstructionClass c = schedule_to_class[i]; - const HloInstruction* instruction = schedule[i]; - - if (concat_bitcast_index == -1 && - c == InstructionClass::kRelatedConcatBitcast) { - concat_bitcast_index = i; - continue; - } - if (concat_bitcast_index != -1 && - c == InstructionClass::kRelatedSliceDone) { - return FailedPrecondition( - "%s", absl::StrCat("Unexpectedly, found concat-bitcast ", - schedule[concat_bitcast_index]->name(), " at ", - concat_bitcast_index, - ", which is before the slice done ", - instruction->name(), " at ", i, ".")); - } - } - - return absl::OkStatus(); - } - - // Return an OK status iff: - // - concat_bitcast and all of its slices come after - // slices_start_after_instruction_name in the schedule AND - // - at least one slice start comes after slices_start_after_instruction_name - // in the schedule, with no non-copy-like instruction between AND - // - if expect_slices_started_at_different_times is true, at least one - // non-copy-like instruction comes between each slice start AND - // - concat_bitcast and all of its slices come before - // slices_done_before_instruction_name in the schedule AND - // - concat_bitcast and all of its slice dones come before - // slices_done_before_instruction_name in the schedule, with no - // non-copy-like instruction between AND - // - concat_bitcast comes after all slice dones AND - static absl::Status CheckSchedule( - const HloModule& module, const HloInstruction* concat_bitcast, - absl::string_view slices_start_after_instruction_name, - absl::string_view slices_done_before_instruction_name, - bool expect_slices_started_at_different_times) { - CHECK(concat_bitcast->IsCustomCall(kConcatBitcastCustomCall)); - - // Get the schedule. - auto entry_schedule = - module.schedule().sequence(module.entry_computation()).instructions(); - - // Initialize schedule_to_class to classify instructions in the schedule. - std::vector schedule_to_class( - entry_schedule.size(), InstructionClass::kUnrelatedNonCopy); - for (int i = 0; i < entry_schedule.size(); ++i) { - const HloInstruction* instruction = entry_schedule[i]; - if (IsAsyncCopyStart(instruction) || IsAsyncCopyDone(instruction) || - IsAsyncSliceStart(instruction) || IsAsyncSliceDone(instruction) || - IsConcatBitcast(instruction)) { - schedule_to_class[i] = InstructionClass::kUnrelatedCopyLike; - } - } - - // Update schedule_to_class with the instructions we care about. - int slices_start_after_index; - ABSL_ASSIGN_OR_RETURN(slices_start_after_index, - FindScheduleIndexOfInstruction( - entry_schedule, slices_start_after_instruction_name, - InstructionClass::kStartAfterNonCopy)); - schedule_to_class[slices_start_after_index] = - InstructionClass::kStartAfterNonCopy; - int slices_done_before_index; - ABSL_ASSIGN_OR_RETURN(slices_done_before_index, - FindScheduleIndexOfInstruction( - entry_schedule, slices_done_before_instruction_name, - InstructionClass::kDoneBeforeNonCopy)); - schedule_to_class[slices_done_before_index] = - InstructionClass::kDoneBeforeNonCopy; - int concat_bitcast_index; - ABSL_ASSIGN_OR_RETURN(concat_bitcast_index, - FindScheduleIndexOfInstruction( - entry_schedule, concat_bitcast->name(), - InstructionClass::kRelatedConcatBitcast)); - schedule_to_class[concat_bitcast_index] = - InstructionClass::kRelatedConcatBitcast; - for (const HloInstruction* slice : concat_bitcast->operands()) { - int done_index; - ABSL_ASSIGN_OR_RETURN(done_index, FindScheduleIndexOfInstruction( - entry_schedule, slice->name(), - InstructionClass::kRelatedSliceDone)); - schedule_to_class[done_index] = InstructionClass::kRelatedSliceDone; - int start_index; - ABSL_ASSIGN_OR_RETURN(start_index, - FindScheduleIndexOfInstruction( - entry_schedule, slice->operand(0)->name(), - InstructionClass::kRelatedSliceStart)); - schedule_to_class[start_index] = InstructionClass::kRelatedSliceStart; - } - - // Perform scheduling checks. - ABSL_RETURN_IF_ERROR(ConcatBitcastAndSlicesAfterInstruction( - entry_schedule, schedule_to_class, slices_start_after_index)); - ABSL_RETURN_IF_ERROR(OneSliceStartAfterInstructionWithNoCopyLikeBetween( - entry_schedule, schedule_to_class, slices_start_after_index)); - if (expect_slices_started_at_different_times) { - ABSL_RETURN_IF_ERROR(AtLeastOneNonCopyLikeInstructionBetweenSliceStarts( - entry_schedule, schedule_to_class)); - } - ABSL_RETURN_IF_ERROR(ConcatBitcastAndSlicesBeforeInstruction( - entry_schedule, schedule_to_class, slices_done_before_index)); - ABSL_RETURN_IF_ERROR( - ConcatBitcastAndSliceDonesBeforeInstructionWithNoCopyLikeBetween( - entry_schedule, schedule_to_class, slices_done_before_index)); - ABSL_RETURN_IF_ERROR( - ConcatBitcastAfterSliceDones(entry_schedule, schedule_to_class)); - - return absl::OkStatus(); - } - - // Returns absl::OkStatus iff: - // - Each slice is assigned a chunk that is the same size as the slice - // instruction's shape. - // - When the slices of sliced_copy_result are sorted in expected spatial - // order, they are assigned chunks that spatially fall in the same order AND - // - The slices of sliced_copy_result are assigned contiguous memory chunks - // AND - // - The sliced_copy_result is assigned a chunk that is the concatenation of - // the slice chunks AND - // - The size of the chunk assigned to the sliced_copy_result has the same - // size as the instruction's shape - static absl::Status CheckSliceChunks(const PresetAssignments& assignments, - const HloInstruction* sliced_copy_result, - bool expect_bitcasted_io = false) { - const HloInstruction* concat_bitcast = - (expect_bitcasted_io ? sliced_copy_result->operand(0) - : sliced_copy_result); - CHECK(concat_bitcast->IsCustomCall(kConcatBitcastCustomCall)); - - absl::flat_hash_map slices_to_chunks; - std::optional result_chunk = std::nullopt; - - for (const auto& position_chunk_pair : assignments.chunks()) { - if (position_chunk_pair.first.instruction == sliced_copy_result) { - if (result_chunk.has_value()) { - return FailedPrecondition( - "%s", absl::StrCat("Sliced copy ", sliced_copy_result->name(), - " is assigned more than one chunk: ", - result_chunk->ToString(), " and ", - position_chunk_pair.second.ToString())); - } - result_chunk = position_chunk_pair.second; - } - for (const HloInstruction* slice : concat_bitcast->operands()) { - if (position_chunk_pair.first.instruction == slice) { - auto it = slices_to_chunks.find(slice); - if (it != slices_to_chunks.end()) { - return FailedPrecondition( - "%s", absl::StrCat("Slice ", slice->name(), - " is assigned more than one chunk: ", - it->second.ToString(), " and ", - position_chunk_pair.second.ToString())); - } - slices_to_chunks[slice] = position_chunk_pair.second; - } - } - } - - std::vector sorted_slices = - SortSlicesInExpectedSpatialOrder(concat_bitcast); - VLOG(1) << "Chunk assignments for " << sliced_copy_result->name() << ":\n" - << absl::StrJoin( - sorted_slices, "\n", - [&](std::string* out, const HloInstruction* slice) { - auto it = slices_to_chunks.find(slice); - std::string chunk = "no chunk assigned"; - if (it != slices_to_chunks.end()) { - chunk = it->second.ToString(); - } - absl::StrAppend(out, " slice ", slice->name(), ": ", - chunk); - }) - << "\n sliced copy result " << sliced_copy_result->name() << ": " - << (result_chunk.has_value() ? result_chunk->ToString() - : "no chunk assigned"); - if (sorted_slices.empty()) { - return absl::OkStatus(); - } - - // Check that slices are assigned contiguous chunks that are spatially - // ordered according to sorted_slices. Also make sure that slices are - // assigned chunks with sizes that match their shape. - int64_t previous_end = -1; - int64_t min_offset = std::numeric_limits::max(); - int64_t max_limit = std::numeric_limits::min(); - for (const HloInstruction* slice : sorted_slices) { - auto it = slices_to_chunks.find(slice); - if (it == slices_to_chunks.end()) { - return FailedPrecondition( - "%s", - absl::StrCat("Slice ", slice->name(), " is not assigned a chunk")); - } - const Chunk& chunk = it->second; - - if (chunk.size != ShapeSize(slice->shape())) { - return FailedPrecondition( - "%s", - absl::StrCat("Slice ", slice->name(), " is assigned chunk ", - chunk.ToString(), " with size ", chunk.size, - ". Expected a size of ", ShapeSize(slice->shape()), - ", to match its shape.")); - } - - if (previous_end != -1 && chunk.offset != previous_end) { - return FailedPrecondition( - "%s", absl::StrCat( - "Slice ", slice->name(), " starts at offset ", - chunk.offset, ". Expected it to start at ", previous_end, - " because that's where the previous slice ended.")); - } - previous_end = chunk.chunk_end(); - min_offset = std::min(min_offset, chunk.offset); - max_limit = std::max(max_limit, chunk.chunk_end()); - } - - // Check that the sliced copy result is assigned a chunk that is the - // concatenation of the slice chunks. - if (!result_chunk.has_value()) { - return FailedPrecondition( - "%s", absl::StrCat("Sliced copy result ", sliced_copy_result->name(), - " is not assigned a chunk.")); - } - Chunk expected_result_chunk = Chunk::FromOffsetEnd(min_offset, max_limit); - if (!(*result_chunk == expected_result_chunk)) { - return FailedPrecondition( - "%s", absl::StrCat("Sliced copy result ", sliced_copy_result->name(), - " is assigned chunk ", result_chunk->ToString(), - ", but it's expected to be assigned chunk ", - expected_result_chunk.ToString())); - } - if (result_chunk->size != ShapeSize(sliced_copy_result->shape())) { - return FailedPrecondition( - "%s", absl::StrCat("Sliced copy result ", sliced_copy_result->name(), - " is assigned chunk ", result_chunk->ToString(), - " with size ", result_chunk->size, - ". Expected a size of ", - ShapeSize(sliced_copy_result->shape()), - ", to match its shape.")); - } - - return absl::OkStatus(); - } - - SlicedPrefetchTest() { - // Force tests to fail if ProposeSlices is unexpectedly called. - EXPECT_CALL(slice_proposer_, ProposeSlices(_, _)).Times(0); - options_ = MakeDefaultOptions(); - } - - // Optional method to setup common ProposeSlices expectations for several - // tests. - void SetupProposeSlicesToExpect2SlicesOfF32x8x8() { - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_8_8_, EqualsSlicedPrefetchOptions( - options().sliced_prefetch_options))) - .WillRepeatedly(Return(SliceProposalCollection({ - SliceProposal({f32_4_8_, std::vector({{0, 4}, {0, 8}}), - ShapeSize(f32_4_8_)}), - SliceProposal({f32_4_8_, std::vector({{4, 8}, {0, 8}}), - ShapeSize(f32_4_8_)}), - }))); - } - - Options MakeDefaultOptions() { - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 1024; - options.sliced_prefetch_options.set_max_slices(2); - options.sliced_prefetch_options.set_min_bytes(8); - options.propose_slice_fn = [&](const Shape& shape, - const SlicedPrefetchOptions& options) { - return slice_proposer_.ProposeSlices(shape, options); - }; - options.get_equivalent_s8_shape_fn = [](const Shape& original_shape) { - return ShapeUtil::MakeShape(S8, {ShapeSize(original_shape)}); - }; - return options; - } - - const Options& options() const { return options_; } - - const Shape f32_8_8_ = ShapeUtil::MakeShape(F32, {8, 8}); - const Shape f32_4_8_ = ShapeUtil::MakeShape(F32, {4, 8}); - MockSliceProposer slice_proposer_; - - private: - Options options_; -}; - -TEST_F(SlicedPrefetchTest, TwoSlices) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - ROOT r = f32[8,8] add(c, p1) -})hlo"; - - SetupProposeSlicesToExpect2SlicesOfF32x8x8(); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), MakeDefaultOptions(), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - auto root = module->entry_computation()->root_instruction(); - - // Expect p1 to be copied via a sliced prefetch for use in r. - EXPECT_THAT(root, op::Add(_, IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 4}, {0, 8}}, {{4, 8}, {0, 8}}}, - op::Parameter(1)))); - - // Check the instruction schedule. - EXPECT_OK(CheckSchedule(*module, root->operand(1), - /*slices_start_after_instruction_name=*/"p1", - /*slices_done_before_instruction_name=*/"r", - /*expect_slices_started_at_different_times=*/true)); - - // Check expectations on the chunks assigned to the asynchronous sliced copy. - EXPECT_OK(CheckSliceChunks(*assignments, root->operand(1))); -} - -// This module is optmized as below, which adds two slices followed by a concat -// bitcast. Since the concat bitcast uses the same slice buffers for its output, -// the heap simulator trace should not have freed any of the slices before the -// concat bitcast users are allocated/processed. -// -// ENTRY main { -// ... -// p1 = f32[8,8]{1,0} parameter(1) -// slice-start = slice-start(p1), slice={[4:8], [0:8]} -// ... -// slice-start.1 = slice-start(p1), slice={[0:4], [0:8]} -// ... -// c = f32[8,8]{1,0:S(1)} tanh(b) -// slice-done = f32[4,8]{1,0:S(1)} slice-done(slice-start) -// slice-done.1 = f32[4,8]{1,0:S(1)} slice-done(slice-start.1) -// custom-call = f32[8,8]{1,0:S(1)} custom-call(slice-done, slice-done.1), -// custom_call_target="ConcatBitcast" -// r = f32[8,8]{1,0:S(1)} add(c, custom-call) -// n = f32[8,8]{1,0:S(1)} negate(custom-call) -// ROOT f = f32[8,8]{1,0} add(r, n) -// } -// -TEST_F(SlicedPrefetchTest, SlicedPrefetchHeapSimulatorTrace) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - r = f32[8,8] add(c, p1) - n = f32[8,8] negate(p1) - ROOT f = f32[8,8] add(r, n) -})hlo"; - - SetupProposeSlicesToExpect2SlicesOfF32x8x8(); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), MakeDefaultOptions(), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - auto* r_instr = FindInstruction(module.get(), "r"); - EXPECT_NE(r_instr, nullptr); - auto* n_instr = FindInstruction(module.get(), "n"); - EXPECT_NE(n_instr, nullptr); - - // Expect p1 to be copied via a sliced prefetch for use in r and n. - EXPECT_THAT( - r_instr, - op::Add(_, IsAsyncSlicedCopy(kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 4}, {0, 8}}, {{4, 8}, {0, 8}}}, - op::Parameter(1)))); - EXPECT_THAT(n_instr, op::Negate(r_instr->operand(1))); - - // Check the instruction schedule. - EXPECT_OK(CheckSchedule(*module, r_instr->operand(1), - /*slices_start_after_instruction_name=*/"p1", - /*slices_done_before_instruction_name=*/"r", - /*expect_slices_started_at_different_times=*/true)); - - // Check expectations on the chunks assigned to the asynchronous sliced copy. - EXPECT_OK(CheckSliceChunks(*assignments, r_instr->operand(1))); - - const HeapSimulatorTrace& heap_trace = - assignments->assignment_information_for_space(kAlternateMemorySpace) - ->heap_simulator_trace; - // Track the set of instructions currently living in the alternate memory - // space. - // - ALLOC event: Instruction should not be in the set. Add it. - // - FREE event: Instruction should be in the set. Remove it. - // - Concat bitcast instruction: Slice operands should remain in the set until - // all concat bitcast users are allocated. - absl::flat_hash_set allocated_instructions; - for (const auto& event : heap_trace.events()) { - VLOG(3) << "event: " << event.DebugString(); - const HloInstruction* instruction = - FindInstruction(module.get(), event.instruction_name()); - EXPECT_NE(instruction, nullptr) - << "Instruction not found: " << event.instruction_name(); - if (instruction->opcode() == HloOpcode::kCustomCall) { - EXPECT_NE(instruction->custom_call_target(), - memory_space_assignment::kConcatBitcastCustomCall) - << "We do not expect concat bitcast custom call to add any " - "independent events to the heap trace."; - } - if (instruction->opcode() == HloOpcode::kSlice) { - EXPECT_TRUE(event.kind() == HeapSimulatorTrace::Event::ALLOC || - event.kind() == HeapSimulatorTrace::Event::FREE); - } - if (event.kind() == HeapSimulatorTrace::Event::ALLOC) { - allocated_instructions.insert(instruction); - } else if (event.kind() == HeapSimulatorTrace::Event::FREE) { - EXPECT_TRUE(allocated_instructions.contains(instruction)) - << "FREE is called on slice instruction before its ALLOC or is its " - "being called more than once on the same slice buffer."; - allocated_instructions.erase(instruction); - } else { - FAIL() << "Unexpected event kind: " << event.kind() - << " for instruction: " << event.instruction_name(); - } - // At the time we allocate the r and n instructions, we should still have - // valid allocations for both slices of the concatbitcast operands, because - // the concatbitcast should share that buffer with its users, i.e. r and n - // instructions. - if ((instruction == r_instr || instruction == n_instr) && - event.kind() == HeapSimulatorTrace::Event::ALLOC) { - int slice_count = absl::c_count_if( - allocated_instructions, [](const HloInstruction* inst) { - return inst->opcode() == HloOpcode::kSlice; - }); - EXPECT_EQ(slice_count, 2) - << "Did not find enough valid allocations for both slice buffers in " - "the trace at the time of allocation for r or n instructions."; - } - } -} - -TEST_F(SlicedPrefetchTest, TwoSlicesWithCopyReplacement) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - p1_copy1 = f32[8,8] copy(p1) - p1_copy2 = f32[8,8] copy(p1) - - r1 = f32[8,8] add(c, p1_copy1) - r2 = f32[8,8] add(c, p1_copy2) - - ROOT r = f32[8,8] add(r1, r2) -})hlo"; - Options options = MakeDefaultOptions(); - options.enable_sync_copy_replacement = true; - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - SetupProposeSlicesToExpect2SlicesOfF32x8x8(); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), std::move(options), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); -} - -TEST_F(SlicedPrefetchTest, ThreeSlices) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - ROOT r = f32[8,8] add(c, p1) -})hlo"; - const Shape f32_3_8 = ShapeUtil::MakeShape(F32, {3, 8}); - const Shape f32_2_8 = ShapeUtil::MakeShape(F32, {2, 8}); - - Options options = MakeDefaultOptions(); - options.sliced_prefetch_options.set_max_slices(3); - - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_8_8_, EqualsSlicedPrefetchOptions( - options.sliced_prefetch_options))) - .WillRepeatedly(Return(SliceProposalCollection({ - SliceProposal({f32_3_8, std::vector({{0, 3}, {0, 8}}), - ShapeSize(f32_3_8)}), - SliceProposal({f32_3_8, std::vector({{3, 6}, {0, 8}}), - ShapeSize(f32_3_8)}), - SliceProposal({f32_2_8, std::vector({{6, 8}, {0, 8}}), - ShapeSize(f32_2_8)}), - }))); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), std::move(options), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - auto root = module->entry_computation()->root_instruction(); - - // Expect p1 to be copied via a sliced prefetch for use in r. - EXPECT_THAT( - root, - op::Add(_, IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 3}, {0, 8}}, {{3, 6}, {0, 8}}, {{6, 8}, {0, 8}}}, - op::Parameter(1)))); - - // Check the instruction schedule. - EXPECT_OK(CheckSchedule(*module, root->operand(1), - /*slices_start_after_instruction_name=*/"p1", - /*slices_done_before_instruction_name=*/"r", - /*expect_slices_started_at_different_times=*/true)); - - // Check expectations on the chunks assigned to the asynchronous sliced copy. - EXPECT_OK(CheckSliceChunks(*assignments, root->operand(1))); -} - -TEST_F(SlicedPrefetchTest, SlicingDisabled) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - ROOT r = f32[8,8] add(c, p1) -})hlo"; - - Options options = MakeDefaultOptions(); - options.sliced_prefetch_options.set_max_slices(0); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), std::move(options), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - // Check that there are not any sliced prefetches in the schedule. - auto entry_schedule = - module->schedule().sequence(module->entry_computation()).instructions(); - for (const HloInstruction* instruction : entry_schedule) { - EXPECT_FALSE(IsAsyncSliceStart(instruction)); - EXPECT_FALSE(IsAsyncSliceDone(instruction)); - EXPECT_FALSE(IsConcatBitcast(instruction)); - } -} - -TEST_F(SlicedPrefetchTest, TooSmallToSlice) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - ROOT r = f32[8,8] add(c, p1) -})hlo"; - - Options options = MakeDefaultOptions(); - options.sliced_prefetch_options.set_min_bytes(1000000000); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), std::move(options), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - // No tensor is big enough to be sliced, so check that there are not any - // sliced prefetches. - auto entry_schedule = - module->schedule().sequence(module->entry_computation()).instructions(); - for (const HloInstruction* instruction : entry_schedule) { - EXPECT_FALSE(IsAsyncSliceStart(instruction)); - EXPECT_FALSE(IsAsyncSliceDone(instruction)); - EXPECT_FALSE(IsConcatBitcast(instruction)); - } -} - -TEST_F(SlicedPrefetchTest, FallbackToUnsliced) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - ROOT r = f32[8,8] add(c, p1) -})hlo"; - - Options options = MakeDefaultOptions(); - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_8_8_, EqualsSlicedPrefetchOptions( - options.sliced_prefetch_options))) - .WillRepeatedly(Return(absl::StatusOr( - FailedPrecondition("%s", "Cannot slice.")))); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), std::move(options), - /*max_prefetch_interval=*/10, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - // No tensor is big enough to be sliced, so check that there are not any - // sliced prefetches. - auto entry_schedule = - module->schedule().sequence(module->entry_computation()).instructions(); - for (const HloInstruction* instruction : entry_schedule) { - EXPECT_FALSE(IsAsyncSliceStart(instruction)); - EXPECT_FALSE(IsAsyncSliceDone(instruction)); - EXPECT_FALSE(IsConcatBitcast(instruction)); - } -} - -TEST_F(SlicedPrefetchTest, UsingCostAnalysisIntervalPicker) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - ROOT r = f32[8,8] add(c, p1) -})hlo"; - - SetupProposeSlicesToExpect2SlicesOfF32x8x8(); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/MakeDefaultOptions()); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - auto root = module->entry_computation()->root_instruction(); - - // Expect p1 to be copied via a sliced prefetch for use in r. - EXPECT_THAT(root, op::Add(_, IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 4}, {0, 8}}, {{4, 8}, {0, 8}}}, - op::Parameter(1)))); - - // Check the instruction schedule. - EXPECT_OK(CheckSchedule( - *module, root->operand(1), - // The CostAnalysisPrefetchIntervalPicker does not necessarily pick the - // earliest possible time to start the prefetch. - /*slices_start_after_instruction_name=*/"a", - /*slices_done_before_instruction_name=*/"r", - /*expect_slices_started_at_different_times=*/true)); - - // Check expectations on the chunks assigned to the asynchronous sliced copy. - EXPECT_OK(CheckSliceChunks(*assignments, root->operand(1))); -} - -TEST_F(SlicedPrefetchTest, LoopAliasing) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -WhileBody { - body_param = (f32[8,8], f32[8,8], f32[], f32[]) parameter(0) - v0 = f32[8,8] get-tuple-element(body_param), index=0 - v1 = f32[8,8] get-tuple-element(body_param), index=1 - i = f32[] get-tuple-element(body_param), index=2 - limit = f32[] get-tuple-element(body_param), index=3 - one = f32[] constant(1) - - new_i = f32[] add(i, one) - new_v1 = f32[8,8] add(v0, v1) - - ROOT while_result = (f32[8,8], f32[8,8], f32[], f32[]) tuple(v0, new_v1, new_i, limit) -} - -WhileCond { - cond_param = (f32[8,8], f32[8,8], f32[], f32[]) parameter(0) - i = f32[] get-tuple-element(cond_param), index=2 - limit = f32[] get-tuple-element(cond_param), index=3 - - ROOT cond_result = pred[] compare(i, limit), direction=LT -} - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - iterations = f32[] parameter(2) - initial = f32[] constant(0) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - - t = (f32[8,8], f32[8,8], f32[], f32[]) tuple(p0, p1, initial, iterations) - w = (f32[8,8], f32[8,8], f32[], f32[]) while(t), condition=WhileCond, body=WhileBody - d = f32[8,8] get-tuple-element(w), index=1 - - ROOT r = f32[8,8] add(c, d) -})hlo"; - - SetupProposeSlicesToExpect2SlicesOfF32x8x8(); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = - AssignMemorySpaceUsingCostAnalysis( - module.get(), /*memory_space_options_override=*/MakeDefaultOptions()); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - auto root = module->entry_computation()->root_instruction(); - - // Expect p1 to be copied with a slice. - ASSERT_THAT( - root, // - // r from module - op::Add(_, - // d from module - op::GetTupleElement( - // w from module - op::While( - // t from module - op::Tuple(_, - IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 4}, {0, 8}}, {{4, 8}, {0, 8}}}, - op::Parameter(1)), - _, _)), - 1))); - - // In the resultant code, ensure that the prefetch of p1 is aliased throughout - // the while loop. - HloInstruction* w = root->mutable_operand(1)->mutable_operand(0); - HloInstruction* t = w->mutable_operand(0); - HloInstruction* concat_bitcast = t->mutable_operand(1); - HloComputation* while_body = w->while_body(); - HloInstruction* body_param = while_body->parameter_instruction(0); - HloComputation* while_cond = w->while_condition(); - HloInstruction* cond_param = while_cond->parameter_instruction(0); - - // Things we expect to alias with the concat_bitcast. - absl::flat_hash_set expected_aliases({ - HloPosition{concat_bitcast, {}}, - HloPosition{w, {1}}, - HloPosition{t, {1}}, - HloPosition{body_param, {1}}, - HloPosition{cond_param, {1}}, - }); - - // Check the aliasing. - auto alias_analysis = - HloAliasAnalysis::Run(module.get(), &alias_info_).value(); - VLOG(2) << alias_analysis->ToString(); - const HloBuffer& concat_bitcast_buffer = - alias_analysis->GetUniqueBufferAt(concat_bitcast); - EXPECT_THAT(concat_bitcast_buffer.ComputePositions(), - ::testing::IsSupersetOf(expected_aliases)); - - // Since expected_aliases are aliased, we expect only 1 to be assigned a - // chunk. - int num_chunks_for_expected_aliases = 0; - for (const auto& position_chunk_pair : assignments->chunks()) { - if (expected_aliases.contains(position_chunk_pair.first)) { - num_chunks_for_expected_aliases++; - } - } - EXPECT_EQ(num_chunks_for_expected_aliases, 1); -} - -class MockRepacker : public MemorySpaceAssignmentRepacker { - public: - MockRepacker() - : MemorySpaceAssignmentRepacker(std::numeric_limits::max(), 1) {} - - MOCK_METHOD(absl::StatusOr, Repack, (absl::Span), - (override)); -}; - -// Here, we test the following things: -// - Without repacking, we are unable to prefetch p4. -// - With repacking, we are able to prefetch p4. -// - When repacking occurs, we expect p2 and p3 to have been allocated chunks. -// We are only proposing slices for f32[32, 16] and not f32[16,16]; thus, we -// expect slicing metadata to be attached to the repacking block for p2 but -// not p3. -// - We make the repacker assign the first slice (in time) of p2 the larger -// offset. After MSA, we check to make sure the fist slice is using the -// larger slicing parameters -TEST_F(SlicedPrefetchTest, Repack) { - absl::string_view hlo_string = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - /* parameters */ - p0 = f32[] parameter(0) - p1 = f32[16,16] parameter(1) - p2 = f32[32,16] parameter(2) - p3 = f32[16,16] parameter(3) - p4 = f32[32,16] parameter(4) - - /* filler that we can prefetch over */ - x1 = f32[] add(p0,p0) - x2 = f32[] add(x1, x1) - - /* uses of p1 and p3 */ - a = f32[16,16] sine(p1) - c = f32[16,16] sine(p3) - - /* more filler, giving us time to prefetch p4, when repacking */ - x3 = f32[] add(x2, x2) - x4 = f32[] add(x3, x3) - - /* uses of p2 and p4 */ - b = f32[32,16] sine(p2) - d = f32[32,16] sine(p4) - - /* make sure that x4, a, b, c, d are not dead code */ - z1 = f32[16,16] broadcast(x4), dimensions={} - z2 = f32[16,16] add(z1, a) - z3 = f32[32,16] concatenate(z2, c), dimensions={0} - z4 = f32[32,16] add(z3, b) - ROOT z5 = f32[32,16] add(z4, d) -})hlo"; - - // Create 2 copies of the module, one to run without repacking and one to run - // with repacking. - ASSERT_OK_AND_ASSIGN(auto module_no_repacking, - ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(auto module_with_repacking, - ParseAndReturnVerifiedModule(hlo_string)); - VLOG(1) << "Original module:\n" - << module_no_repacking->ToString(HloPrintOptions::ShortParsable()); - - // Setup slicing expectations so that we slice f32[32, 16], but not - // f32[16,16]. - Shape f32_16_16 = ShapeUtil::MakeShape(F32, {16, 16}); - Shape f32_32_16 = ShapeUtil::MakeShape(F32, {32, 16}); - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_16_16, EqualsSlicedPrefetchOptions( - options().sliced_prefetch_options))) - .WillRepeatedly(Return(SliceProposalCollection({}))); - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_32_16, EqualsSlicedPrefetchOptions( - options().sliced_prefetch_options))) - .WillRepeatedly(Return(SliceProposalCollection({ - SliceProposal({f32_16_16, std::vector({{0, 16}, {0, 16}}), - ShapeSize(f32_16_16)}), - SliceProposal({f32_16_16, - std::vector({{16, 32}, {0, 16}}), - ShapeSize(f32_16_16)}), - }))); - - // Force MSA to prefer prefetching (in order) p1, p2, p3, p4, and then - // anything else. - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that for p1, p2, p3, and - // p4. - int priority = 100; - if (x.buffer->instruction()->name() == "p1") { - priority = 1; - } else if (x.buffer->instruction()->name() == "p2") { - priority = 2; - } else if (x.buffer->instruction()->name() == "p3") { - priority = 3; - } else if (x.buffer->instruction()->name() == "p4") { - priority = 4; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - - // Configure MSA. - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 50); - Options options_no_repacking = MakeDefaultOptions(); - options_no_repacking.max_size_in_bytes = 4 * 1024; - options_no_repacking.max_repacks = 0; - - // Run MSA without repacking - std::unique_ptr assignments = AssignMemorySpace( - module_no_repacking.get(), std::move(options_no_repacking), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Post-MSA module (no repacking):\n" - << module_no_repacking->ToString(HloPrintOptions::ShortParsable()); - - // If repacking is disabled, p4 (the source of d) should not be prefetched. - const HloInstruction* d = - FindNamedScheduledInstruction(*module_no_repacking, "d"); - ASSERT_NE(d, nullptr); - EXPECT_FALSE(IsConcatBitcast(d->operand(0))); - - // Configure MSA to repack. - MockRepacker repacker; - absl::flat_hash_map, int64_t> repack_map; - EXPECT_CALL(repacker, Repack(_)) - .WillRepeatedly([](absl::Span allocations) - -> absl::StatusOr { - bool found_p2 = false; - bool found_p3 = false; - for (AllocationBlock* block : allocations) { - VLOG(1) << "Allocation block: " << block->ToString(); - - if (block->inclusive_start_time == 3 && - block->initial_offset == 1024 && block->size == 2048) { - // Move "p2" from offset 1024 -> 2048. - found_p2 = true; - block->offset = 2048; - // We expect p2 to be sliced. Check that it has slicing information - // in its AllocationBlock. - EXPECT_TRUE(block->original_slice_data.has_value()); - if (block->original_slice_data.has_value()) { - SlicedAllocationData expected( - {{AllocatedSlice{1024, 1024, /*inclusive_start_time=*/3}, - AllocatedSlice{1024, 2048, /*inclusive_start_time=*/7}}}); - EXPECT_EQ(*block->original_slice_data, expected) - << "\nExpected: " << expected.ToString() - << "\nGot: " << block->original_slice_data->ToString(); - // Set the first slice for p2 to be place at the larger offset. - block->repacked_slice_data = SlicedAllocationData( - {{AllocatedSlice{1024, 2048, /*inclusive_start_time=*/7}, - AllocatedSlice{1024, 3072, /*inclusive_start_time=*/3}}}); - } - } else if (block->inclusive_start_time == 4 && - block->initial_offset == 3072 && block->size == 1024) { - // Move "p3" from offset 3072 -> 1024. - found_p3 = true; - block->offset = 1024; - // We do not expect p3 to be sliced. Thus, it should not have - // slicing information in its AllocationBlock. - EXPECT_FALSE(block->original_slice_data.has_value()); - } else { - block->offset = block->initial_offset; - } - } - - EXPECT_TRUE(found_p2); - EXPECT_TRUE(found_p3); - - return true; - }); - Options options_with_repacking = MakeDefaultOptions(); - options_with_repacking.max_size_in_bytes = 4 * 1024; - options_with_repacking.max_repacks = 0; - options_with_repacking.max_repacks = 1; - options_with_repacking.repacker = &repacker; - assignments = AssignMemorySpace( - module_with_repacking.get(), std::move(options_with_repacking), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Post-MSA module (with repacking):\n" - << module_with_repacking->ToString(HloPrintOptions::ShortParsable()); - - // If repacking is enabled, p4 (the source of d) should be prefetched. - d = FindNamedScheduledInstruction(*module_with_repacking, "d"); - ASSERT_NE(d, nullptr); - EXPECT_TRUE(IsConcatBitcast(d->operand(0))); - - // Check expectations on the chunks assigned to the asynchronous sliced copy. - // In particular, we want to make sure the slices are still contiguous. - EXPECT_OK(CheckSliceChunks(*assignments, d->operand(0))); - - // Find the slices and offsets for p2, in the order they start in the - // schedule. - std::vector p2_slice_dones; - for (const HloInstruction* i : - module_with_repacking->entry_computation()->instructions()) { - if (IsAsyncSliceStart(i) && i->operand_count() == 1 && - i->operand(0)->name() == "p2") { - ASSERT_EQ(i->user_count(), 1); - p2_slice_dones.push_back(i->users()[0]); - } - } - ASSERT_EQ(p2_slice_dones.size(), 2); - std::vector p2_slice_offsets; - for (const HloInstruction* i : p2_slice_dones) { - for (const auto& position_chunk_pair : assignments->chunks()) { - if (position_chunk_pair.first.instruction == i) { - p2_slice_offsets.push_back(position_chunk_pair.second.offset); - } - } - } - ASSERT_EQ(p2_slice_offsets.size(), 2); - - // We expect the first slice of p2 to have the larger offsets. - EXPECT_THAT(p2_slice_dones[0]->async_wrapped_instruction()->slice_starts(), - ::testing::ElementsAreArray({16, 0})); - EXPECT_THAT(p2_slice_dones[0]->async_wrapped_instruction()->slice_limits(), - ::testing::ElementsAreArray({32, 16})); - EXPECT_EQ(p2_slice_offsets[0], 3072); - EXPECT_THAT(p2_slice_dones[1]->async_wrapped_instruction()->slice_starts(), - ::testing::ElementsAreArray({0, 0})); - EXPECT_THAT(p2_slice_dones[1]->async_wrapped_instruction()->slice_limits(), - ::testing::ElementsAreArray({16, 16})); - EXPECT_EQ(p2_slice_offsets[1], 2048); -} - -struct ModuleAndAssignments { - std::unique_ptr module; - std::unique_ptr assignments; -}; - -// In this test, we ensure that sliced prefetching does not attempt to start a -// slice during a different computation than the one where the slice finishes. -// We do this by forcing a sliced prefetch to start just before back-to-back -// while loops and to immediately finish after them. We use while loops with -// different expected elapse times, so that the ideal place to start the second -// slice is during one of the while loops. -TEST_F(SlicedPrefetchTest, BackToBackWhileLoops) { - // Define constants for building our test HLO. - const std::string while_cond = R"hlo( -WhileCond$ID { - cond_param = (f32[8,8], f32[8,8], f32[], f32[]) parameter(0) - i = f32[] get-tuple-element(cond_param), index=2 - limit = f32[] get-tuple-element(cond_param), index=3 - - ROOT cond_result = pred[] compare(i, limit), direction=LT -})hlo"; - - const std::string while_body = R"hlo( -WhileBody$ID { - body_param = (f32[8,8], f32[8,8], f32[], f32[]) parameter(0) - v0 = f32[8,8] get-tuple-element(body_param), index=0 - v1 = f32[8,8] get-tuple-element(body_param), index=1 - i = f32[] get-tuple-element(body_param), index=2 - limit = f32[] get-tuple-element(body_param), index=3 - one = f32[] constant(1) - - new_i = f32[] add(i, one) - $COMPUTATION - - ROOT while_result = (f32[8,8], f32[8,8], f32[], f32[]) tuple(v0, new_v1, new_i, limit) -})hlo"; - - const std::string while_computation_cheap = R"hlo( - new_v1 = f32[8,8] add(v0, v1))hlo"; - - std::string while_computation_expensive = R"hlo( - new_v1_0 = f32[8,8] add(v0, v1) - new_v1_1 = f32[8,8] tanh(new_v1_0) - new_v1_2 = f32[8,8] tanh(new_v1_1) - new_v1_3 = f32[8,8] tanh(new_v1_2) - new_v1 = f32[8,8] tanh(new_v1_3))hlo"; - - std::string module_text = R"hlo( -HloModule Slice, is_scheduled=true - -$WHILEBODY1 -$WHILECOND1 -$WHILEBODY2 -$WHILECOND2 - -ENTRY main { - loop1_input1 = f32[8,8] parameter(0) - loop1_input2 = f32[8,8] parameter(1) - loop1_iterations = f32[] parameter(2) - loop1_begin = f32[] constant(0) - loop1_tuple = (f32[8,8], f32[8,8], f32[], f32[]) tuple(loop1_input1, loop1_input2, loop1_iterations, loop1_begin) - loop2_input1 = f32[8,8] parameter(3) - loop2_input2 = f32[8,8] parameter(4) - loop2_iterations = f32[] parameter(5) - loop2_begin = f32[] constant(0) - loop2_tuple = (f32[8,8], f32[8,8], f32[], f32[]) tuple(loop2_input1, loop2_input2, loop2_iterations, loop2_begin) - - prefetch = f32[8,8] parameter(6) - loop1_output = (f32[8,8], f32[8,8], f32[], f32[]) while(loop1_tuple), condition=WhileCond1, body=WhileBody1 - loop2_output = (f32[8,8], f32[8,8], f32[], f32[]) while(loop2_tuple), condition=WhileCond2, body=WhileBody2 - prefetch_use = f32[8,8] tanh(prefetch) - - loop1_result = f32[8,8] get-tuple-element(loop1_output), index=1 - loop2_result = f32[8,8] get-tuple-element(loop2_output), index=1 - - tmp1 = f32[8,8] add(loop1_result, loop2_result) - ROOT r = f32[8,8] add(tmp1, prefetch_use) -})hlo"; - - // A lambda for generating HLO with 2 while loops called back to back. The - // first while loop will execute while_computation1 and the second while loop - // will execute while_computation2. - auto gen_hlo = [&](absl::string_view while_computation1, - absl::string_view while_computation2) { - return absl::StrReplaceAll( - module_text, - { - {"$WHILEBODY1", - absl::StrReplaceAll( - while_body, - {{"$ID", "1"}, {"$COMPUTATION", while_computation1}})}, - {"$WHILECOND1", absl::StrReplaceAll(while_cond, {{"$ID", "1"}})}, - {"$WHILEBODY2", - absl::StrReplaceAll( - while_body, - {{"$ID", "2"}, {"$COMPUTATION", while_computation2}})}, - {"$WHILECOND2", absl::StrReplaceAll(while_cond, {{"$ID", "2"}})}, - }); - }; - - // Configure MSA. - SetupProposeSlicesToExpect2SlicesOfF32x8x8(); - // Force MSA to prefer prefetching 'prefetch'. - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that used for 'prefetch'. - int priority = 100; - if (x.buffer->instruction()->name() == "prefetch") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - // We set the minimum prefetch interval to a large enough value (32) to force - // us to prefetch around both while loops, and not just 1. - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(32, 100); - Options options0 = MakeDefaultOptions(); - options0.max_size_in_bytes = 4 * 64; - - // Define a lambda for running MSA on the specified HLO, with the - // configuration above. - auto run_msa = - [&](Options options, - absl::string_view hlo_text) -> absl::StatusOr { - ModuleAndAssignments module_and_assignments; - ABSL_ASSIGN_OR_RETURN(module_and_assignments.module, - ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module_and_assignments.module->ToString( - HloPrintOptions::ShortParsable()); - module_and_assignments.assignments = AssignMemorySpace( - module_and_assignments.module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - VLOG(1) << "Post-MSA module:\n" - << module_and_assignments.module->ToString( - HloPrintOptions::ShortParsable()); - return module_and_assignments; - }; - - // In this case, less time elapses during the first while loop than the - // second. Make sure we start the second slice between the two while loops, - // rather than during the second while loop. - ASSERT_OK_AND_ASSIGN( - ModuleAndAssignments module_and_assignments1, - run_msa(std::move(options0), - gen_hlo(while_computation_cheap, while_computation_expensive))); - auto root1 = - module_and_assignments1.module->entry_computation()->root_instruction(); - EXPECT_THAT(root1, op::Add(_, op::Tanh(IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 4}, {0, 8}}, {{4, 8}, {0, 8}}}, - op::Parameter(6))))); - EXPECT_OK(CheckSchedule( - *module_and_assignments1.module, root1->operand(1)->operand(0), - /*slices_start_after_instruction_name=*/"prefetch", - /*slices_done_before_instruction_name=*/"prefetch_use", - /*expect_slices_started_at_different_times=*/true)); - auto entry_schedule1 = - module_and_assignments1.module->schedule() - .sequence(module_and_assignments1.module->entry_computation()) - .instructions(); - ASSERT_OK_AND_ASSIGN( - std::vector start_indices, - GetSliceStartIndices(entry_schedule1, root1->operand(1)->operand(0))); - ASSERT_EQ(start_indices.size(), 2); - ASSERT_OK_AND_ASSIGN( - int first_while, - FindScheduleIndexOfInstruction( - entry_schedule1, "loop1_output", - SlicedPrefetchTest::InstructionClass::kUnrelatedNonCopy)); - ASSERT_OK_AND_ASSIGN( - int second_while, - FindScheduleIndexOfInstruction( - entry_schedule1, "loop2_output", - SlicedPrefetchTest::InstructionClass::kUnrelatedNonCopy)); - EXPECT_TRUE( - absl::c_is_sorted>( - {start_indices[0], first_while, start_indices[1], second_while}) || - absl::c_is_sorted>( - {start_indices[1], first_while, start_indices[0], second_while})); - - Options options1 = MakeDefaultOptions(); - options1.max_size_in_bytes = 4 * 64; - - // In this case, more time elapses during the first while loop than the - // second. This should push us to use a normal prefetch, rather than slicing, - // since the ideal time to start the second slice will get pushed before - // both while loops. - ASSERT_OK_AND_ASSIGN( - ModuleAndAssignments module_and_assignments2, - run_msa(std::move(options1), - gen_hlo(while_computation_expensive, while_computation_cheap))); - auto root2 = - module_and_assignments2.module->entry_computation()->root_instruction(); - EXPECT_THAT(root2, op::Add(_, op::Tanh(op::AsyncCopy(kAlternateMemorySpace, - kDefaultMemorySpace, - op::Parameter(6))))); - auto entry_schedule2 = - module_and_assignments2.module->schedule() - .sequence(module_and_assignments2.module->entry_computation()) - .instructions(); - ASSERT_OK_AND_ASSIGN( - int copy_done, - FindScheduleIndexOfInstruction( - entry_schedule2, root2->operand(1)->operand(0)->name(), - SlicedPrefetchTest::InstructionClass::kUnrelatedNonCopy)); - ASSERT_OK_AND_ASSIGN( - int copy_start, - FindScheduleIndexOfInstruction( - entry_schedule2, root2->operand(1)->operand(0)->operand(0)->name(), - SlicedPrefetchTest::InstructionClass::kUnrelatedNonCopy)); - ASSERT_OK_AND_ASSIGN( - first_while, - FindScheduleIndexOfInstruction( - entry_schedule2, "loop1_output", - SlicedPrefetchTest::InstructionClass::kUnrelatedNonCopy)); - ASSERT_OK_AND_ASSIGN( - second_while, - FindScheduleIndexOfInstruction( - entry_schedule2, "loop2_output", - SlicedPrefetchTest::InstructionClass::kUnrelatedNonCopy)); - EXPECT_TRUE(absl::c_is_sorted>( - {copy_start, first_while, second_while, copy_done})); -} - -using RepackingTest = ::testing::Test; - -TEST_F(RepackingTest, Colocations) { - AllocationBlock a{10, 20, 100, 0, 1000, 0}; - AllocationBlock b{15, 25, 150, 0, 2000, 1}; - AllocationBlock c{18, 22, 50, 0, 500, 2}; - AllocationBlock d{5, 9, 20, 0, 3000, 3}; - AllocationBlock e{17, 22, 100, 0, 1500, 4}; - AllocationBlock f{25, 27, 150, 0, 2500, 5}; - - // a doesn't have other colocations. - a.next_colocated = &a; - // b and c are colocated. - b.next_colocated = &c; - c.next_colocated = &b; - // d, e, and f are colocated. - d.next_colocated = &f; - e.next_colocated = &d; - f.next_colocated = &e; - - EXPECT_EQ(a.GetColocationsCount(), 1); - EXPECT_THAT(a.GetColocations(), UnorderedElementsAre(&a)); - EXPECT_EQ(b.GetColocationsCount(), 2); - EXPECT_THAT(b.GetColocations(), UnorderedElementsAre(&b, &c)); - EXPECT_EQ(c.GetColocationsCount(), 2); - EXPECT_THAT(c.GetColocations(), UnorderedElementsAre(&b, &c)); - EXPECT_EQ(d.GetColocationsCount(), 3); - EXPECT_THAT(d.GetColocations(), UnorderedElementsAre(&d, &e, &f)); - EXPECT_EQ(e.GetColocationsCount(), 3); - EXPECT_THAT(e.GetColocations(), UnorderedElementsAre(&d, &e, &f)); - EXPECT_EQ(f.GetColocationsCount(), 3); - EXPECT_THAT(f.GetColocations(), UnorderedElementsAre(&d, &e, &f)); -} - -TEST_F(SlicedPrefetchTest, UniformSizedSlicing) { - std::string hlo_text = R"hlo( -HloModule Slice, is_scheduled=true - -ENTRY main { - p0 = f32[8,8] parameter(0) - p1 = f32[8,8] parameter(1) - p2 = f32[8,16] parameter(2) - constant1 = f32[] constant(1.1) - - a = f32[8,8] tanh(p0) - b = f32[8,8] tanh(a) - c = f32[8,8] tanh(b) - d = f32[8,8] tanh(c) - e = f32[8,8] tanh(d) - f = f32[8,8] tanh(e) - g = f32[8,8] tanh(f) - h = f32[8,8] tanh(g) - - x = f32[8,8] add(p1, h) - padded_x = f32[8,16] pad(x, constant1), padding=0_0x0_8 - ROOT r = f32[8,16] add(padded_x, p2) -})hlo"; - const Shape f32_8_16 = ShapeUtil::MakeShape(F32, {8, 16}); - const Shape s8_128 = ShapeUtil::MakeShape(S8, {128}); - - Options options = MakeDefaultOptions(); - options.sliced_prefetch_options.set_max_slices(100000); - options.sliced_prefetch_options.set_preferred_slice_size(4 * 8 * 4); - // TODO(b/563782724): MSA assigns overlapping alternate memory chunks for this - // module. Re-enable verification once that is fixed. - options.verify = false; - - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_8_8_, EqualsSlicedPrefetchOptions( - options.sliced_prefetch_options))) - .WillRepeatedly(Return(SliceProposalCollection({ - SliceProposal( - {s8_128, std::vector({{0, 128}}), ShapeSize(s8_128)}), - SliceProposal({s8_128, std::vector({{128, 256}}), - ShapeSize(s8_128)}), - }))); - - EXPECT_CALL(slice_proposer_, - ProposeSlices(f32_8_16, EqualsSlicedPrefetchOptions( - options.sliced_prefetch_options))) - .WillRepeatedly(Return(SliceProposalCollection({ - SliceProposal( - {s8_128, std::vector({{0, 128}}), ShapeSize(s8_128)}), - SliceProposal({s8_128, std::vector({{128, 256}}), - ShapeSize(s8_128)}), - SliceProposal({s8_128, std::vector({{256, 384}}), - ShapeSize(s8_128)}), - SliceProposal({s8_128, std::vector({{384, 512}}), - ShapeSize(s8_128)}), - }))); - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - VLOG(1) << "Original module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - std::unique_ptr assignments = AssignMemorySpace( - module.get(), std::move(options), - /*max_prefetch_interval=*/100, /*min_prefetch_interval=*/1); - - VLOG(1) << "Post-MSA module:\n" - << module->ToString(HloPrintOptions::ShortParsable()); - - auto root = module->entry_computation()->root_instruction(); - - // Expect p1 to be asynchronously copied via 2 slices, and p2 to be - // asynchronously copied via 4 slices. We expect p1 and p2 to be bitcast - // before slicing and after slicing. - EXPECT_THAT( - root, - op::Add(op::Pad(op::Add(IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 128}}, {{128, 256}}}, op::Parameter(1), - /*expect_bitcasted_io=*/true), - /*don't care*/ _), - /*padding constant*/ _), - IsAsyncSlicedCopy( - kAlternateMemorySpace, kDefaultMemorySpace, - {{{0, 128}}, {{128, 256}}, {{256, 384}}, {{384, 512}}}, - op::Parameter(2), /*expect_bitcasted_io=*/true))); - - // Check expectations on the chunks assigned to the asynchronous sliced copy. - EXPECT_OK(CheckSliceChunks(*assignments, root->operand(1), - /*expect_bitcasted_io=*/true)); - EXPECT_OK(CheckSliceChunks(*assignments, - root->operand(0)->operand(0)->operand(0), - /*expect_bitcasted_io=*/true)); -} - -TEST_F(MemorySpaceAssignmentTest, TestAsyncCopyCustomKernel) { - absl::string_view hlo_string = R"hlo( -HloModule jit_f, is_scheduled=true, entry_computation_layout={(f32[8,128]{1,0:T(8,128)})->(f32[8,128]{1,0:T(8,128)}, f32[8,128]{1,0:T(8,128)})}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true,true} - -ENTRY %main.13 (Arg_0.1: f32[8,128]) -> (f32[8,128], f32[8,128]) { - %Arg_0.1 = f32[8,128]{1,0:T(8,128)} parameter(0), metadata={op_name="x"} - %copy.9 = f32[8,128]{1,0:T(8,128)} copy(%Arg_0.1) - %copy.6 = f32[8,128]{1,0:T(8,128)} copy(%copy.9) - %copy_start.2 = (f32[8,128]{1,0:T(8,128)}, f32[8,128]{1,0:T(8,128)}, s32[]{:T(128)S(2)}) custom-call(%copy.9), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[8,128]{1,0}}, output_to_operand_aliasing={{0}: (0, {})}, control-predecessors={%copy.6} - %get-tuple-element.5 = s32[]{:T(128)S(2)} get-tuple-element(%copy_start.2), index=2 - %get-tuple-element.4 = f32[8,128]{1,0:T(8,128)} get-tuple-element(%copy_start.2), index=1 - %get-tuple-element.3 = f32[8,128]{1,0:T(8,128)} get-tuple-element(%copy_start.2), index=0 - %copy_done.2 = f32[8,128]{1,0:T(8,128)} custom-call(%get-tuple-element.3, %get-tuple-element.4, %get-tuple-element.5), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[8,128]{1,0}, f32[8,128]{1,0}, s32[]}, output_to_operand_aliasing={{}: (1, {})}, control-predecessors={%copy.6} - %copy_start.3 = (f32[8,128]{1,0:T(8,128)}, f32[8,128]{1,0:T(8,128)}, s32[]{:T(128)S(2)}) custom-call(%copy.6), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[8,128]{1,0}}, output_to_operand_aliasing={{0}: (0, {})} - %get-tuple-element.9 = s32[]{:T(128)S(2)} get-tuple-element(%copy_start.3), index=2 - %get-tuple-element.8 = f32[8,128]{1,0:T(8,128)} get-tuple-element(%copy_start.3), index=1 - %get-tuple-element.7 = f32[8,128]{1,0:T(8,128)} get-tuple-element(%copy_start.3), index=0 - %copy_done.3 = f32[8,128]{1,0:T(8,128)} custom-call(%get-tuple-element.7, %get-tuple-element.8, %get-tuple-element.9), custom_call_target="tpu_custom_call", operand_layout_constraints={f32[8,128]{1,0}, f32[8,128]{1,0}, s32[]}, output_to_operand_aliasing={{}: (1, {})} - ROOT %tuple.12 = (f32[8,128]{1,0:T(8,128)}, f32[8,128]{1,0:T(8,128)}) tuple(%copy_done.2, %copy_done.3) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* copy_start_2 = FindInstruction(module.get(), "copy_start.2"); - HloPosition copy_start_2_position{copy_start_2, {1}}; - HloUse copy_start_2_use{copy_start_2, 0, {}}; - HloInstruction* copy_start_3 = FindInstruction(module.get(), "copy_start.3"); - HloPosition copy_start_3_position{copy_start_3, {1}}; - HloUse copy_start_3_use{copy_start_3, 0, {}}; - HloInstruction* copy_done_2 = FindInstruction(module.get(), "copy_done.2"); - HloUse copy_done_2_use{copy_done_2, 1, {}}; - HloInstruction* copy_done_3 = FindInstruction(module.get(), "copy_done.3"); - HloUse copy_done_3_use{copy_done_3, 1, {}}; - - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {copy_start_2_position, kAlternateMemorySpace}, - {copy_start_3_position, kAlternateMemorySpace}, - {copy_done_2_use, kAlternateMemorySpace}, - {copy_done_3_use, kAlternateMemorySpace}, - {copy_start_2_use, kDefaultMemorySpace}, - {copy_start_3_use, kDefaultMemorySpace}}; - - memory_space_options.alternate_memory_space = 1; - memory_space_options.max_size_in_bytes = 50266112; - memory_space_options.alignment_in_bytes = 4096; - memory_space_options.replicated_split_dimension = -1; - memory_space_options.any_split_dimension = -2; - memory_space_options.reduce_scoped_memory_limit = false; - memory_space_options.allocate_reserved_scoped_memory_at_same_offset = true; - memory_space_options.max_outstanding_prefetches = 40; - memory_space_options.max_outstanding_evictions = 40; - memory_space_options.while_use_extra_outstanding_prefetch_limit = 4; - memory_space_options.max_retries = 2; - memory_space_options.max_repacks = 0; - memory_space_options.repack_after_every_allocation = false; - memory_space_options.enable_cross_program_prefetch = true; - memory_space_options.default_cross_program_prefetch_heuristic = false; - memory_space_options.enable_cross_program_prefetch_freeing = true; - memory_space_options.max_cross_program_prefetches = 1; - memory_space_options.cross_program_prefetch_permissive_mode = false; - memory_space_options.enable_while_redundant_eviction_elimination = true; - memory_space_options.use_repeated_instance_for_preferred_prefetch_time = - false; - memory_space_options.enforce_prefetch_fifo_order = false; - memory_space_options.enable_sync_copy_replacement = true; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.extend_async_copies_limit_for_sync_mem_op_conversion = 0; - memory_space_options.inefficient_use_to_copy_ratio = 0.5; - memory_space_options.always_spill_to_default_memory = false; - memory_space_options.window_prefetch_mode = WindowPrefetchMode::kNone; - memory_space_options.expanded_scoped_alternate_memory_mode = - ExpandedScopedAlternateMemoryMode::DISABLED; - memory_space_options.post_module_scoped_alternate_memory_size_in_bytes = 0; - memory_space_options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - return position.instruction->name() == "copy_start.2" || - position.instruction->name() == "copy_start.3"; - }; - memory_space_options.is_allowed_in_alternate_mem_fn = - [&](const HloValue& value) { - // Check if the value belongs in the entry computation. - HloInstruction* instruction = value.instruction(); - HloComputation* computation = instruction->parent(); - if ((instruction->name() == "copy_start.2" || - instruction->name() == "copy_start.3") && - value.index() == ShapeIndex({1})) { - return true; - } - bool in_entry_computation = - (computation == computation->parent()->entry_computation()); - if (in_entry_computation && - instruction->opcode() == HloOpcode::kParameter) { - return value.shape().has_layout() && - value.shape().layout().memory_space() != kDefaultMemorySpace; - } - return true; - }; - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - - EXPECT_EQ(FindInstruction(module.get(), "copy_done.2") - ->operand(1) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "copy_done.3") - ->operand(1) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "copy_done.2") - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - - EXPECT_EQ(FindInstruction(module.get(), "copy_done.3") - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, TestCustomKernelColoringBug) { - absl::string_view hlo_string = R"hlo( -HloModule jit_run_benchmark, is_scheduled=true, entry_computation_layout={(bf16[1024,512]{1,0:T(8,128)(2,1)}, bf16[2,512,4096]{2,1,0:T(8,128)(2,1)}, s32[3]{0:T(128)})->bf16[1024,512]{1,0:T(8,128)(2,1)}}, allow_spmd_sharding_propagation_to_parameters={false,false,false}, allow_spmd_sharding_propagation_to_output={true}, num_partitions=8 - -ENTRY %main.28_spmd (param.1: bf16[1024,512], param.2: bf16[2,512,4096], param: s32[3]) -> bf16[1024,512] { - %constant.1 = s32[1]{0:T(128)} constant({0}), metadata={op_name="call.0"} - %param.2 = bf16[2,512,4096]{2,1,0:T(8,128)(2,1)} parameter(1), sharding={devices=[1,8,1]<=[8]}, metadata={op_name="args[1]"} - %param.1 = bf16[1024,512]{1,0:T(8,128)(2,1)} parameter(0), sharding={devices=[1,8]<=[8]}, metadata={op_name="args[0]"} - %param = s32[3]{0:T(128)} parameter(2), sharding={replicated}, metadata={op_name="args[2]"} - %make_group_metadata.1 = (s32[1]{0:T(128)S(4)}, s32[1]{0:T(128)S(4)}, s32[4]{0:T(128)S(4)}, s32[3]{0:T(128)S(4)}, s32[3]{0:T(128)S(4)}) custom-call(%param, %constant.1), custom_call_target="tpu_custom_call", operand_layout_constraints={s32[3]{0}, s32[1]{0}}, backend_config={"custom_call_config": {"body": "", "serialization_format": 1, "needs_layout_passes": true, "output_memory_space_colors": [{"shape_index":[0],"color":4},{"shape_index":[1],"color":4},{"shape_index":[2],"color":4},{"shape_index":[3],"color":4},{"shape_index":[4],"color":4}]}} - %get-tuple-element.11 = s32[1]{0:T(128)S(4)} get-tuple-element(%make_group_metadata.1), index=1 - %get-tuple-element.10 = s32[1]{0:T(128)S(4)} get-tuple-element(%make_group_metadata.1), index=0 - %get-tuple-element.28 = s32[3]{0:T(128)S(4)} get-tuple-element(%make_group_metadata.1), index=4 - %get-tuple-element.22 = s32[3]{0:T(128)S(4)} get-tuple-element(%make_group_metadata.1), index=3 - %get-tuple-element.16 = s32[4]{0:T(128)S(4)} get-tuple-element(%make_group_metadata.1), index=2 - %ragged_latency_optimized_reduce_scatter_out_feature_matmul.1 = bf16[1024,512]{1,0:T(8,128)(2,1)} custom-call(%get-tuple-element.10, %get-tuple-element.11, %get-tuple-element.16, %get-tuple-element.22, %get-tuple-element.28, /*index=5*/%param.1, %param.2), custom_call_target="tpu_custom_call", operand_layout_constraints={s32[1]{0}, s32[1]{0}, s32[4]{0}, s32[3]{0}, s32[3]{0}, bf16[1024,512]{1,0}, bf16[2,512,4096]{2,1,0}}, backend_config={"flag_configs":[{"flag_type":"XLA_TPU_EVENLY_DISTRIBUTE_CRITICAL_NODES","value":{"boolean_value":true}},{"flag_type":"XLA_TPU_ENABLE_LLO_PIPELINING","value":{"boolean_value":true}},{"flag_type":"XLA_TPU_ENABLE_LLO_SMALL_LOOP_FLATTENING","value":{"boolean_value":true}},{"flag_type":"XLA_TPU_RESCHEDULE_DMA_DONES","value":{"boolean_value":true}},{"flag_type":"XLA_TPU_MAX_INSTS_TO_PREDICATE","value":{"integer_value":"65536"}}],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"custom_call_config":{"body":"","has_communication":true,"collective_id":"1","cost_estimate":{"flops":"5","transcendentals":"0","bytes_accessed":"15000000"},"needs_layout_passes":true,"allow_input_fusion":[],"serialization_format":"1","output_memory_colors":[],"output_memory_space_colors":[{"color":"1","shape_index":[]}],"input_memory_space_colors":[{"operand_index":"5","color":"1","shape_index":[]},{"operand_index":"6","color":"1","shape_index":[]}]},"used_scoped_memory_configs":[]} - ROOT %copy.22 = bf16[1024,512]{1,0:T(8,128)(2,1)} copy(%ragged_latency_optimized_reduce_scatter_out_feature_matmul.1) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* custom_kernel = FindInstruction( - module.get(), - "ragged_latency_optimized_reduce_scatter_out_feature_matmul.1"); - HloPosition custom_kernel_position{custom_kernel, {}}; - HloUse custom_kernel_param_1_use{custom_kernel, 5, {}}; - HloUse custom_kernel_param_2_use{custom_kernel, 6, {}}; - - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {custom_kernel_position, kAlternateMemorySpace}, - {custom_kernel_param_1_use, kAlternateMemorySpace}, - {custom_kernel_param_2_use, kAlternateMemorySpace}}; - - memory_space_options.default_memory_space = 0; - memory_space_options.alternate_memory_space = 1; - memory_space_options.max_size_in_bytes = 62537728; - memory_space_options.alignment_in_bytes = 4096; - memory_space_options.replicated_split_dimension = -1; - memory_space_options.any_split_dimension = -2; - memory_space_options.reduce_scoped_memory_limit = false; - memory_space_options.allocate_reserved_scoped_memory_at_same_offset = true; - memory_space_options.max_outstanding_prefetches = 40; - memory_space_options.max_outstanding_evictions = 40; - memory_space_options.while_use_extra_outstanding_prefetch_limit = 4; - memory_space_options.max_retries = 2; - memory_space_options.max_repacks = 4; - memory_space_options.repack_after_every_allocation = false; - memory_space_options.enable_cross_program_prefetch = true; - memory_space_options.default_cross_program_prefetch_heuristic = false; - memory_space_options.enable_cross_program_prefetch_freeing = true; - memory_space_options.max_cross_program_prefetches = 1; - memory_space_options.cross_program_prefetch_permissive_mode = false; - memory_space_options.enable_while_redundant_eviction_elimination = true; - memory_space_options.use_repeated_instance_for_preferred_prefetch_time = - false; - memory_space_options.enforce_prefetch_fifo_order = false; - memory_space_options.enable_sync_copy_replacement = true; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.extend_async_copies_limit_for_sync_mem_op_conversion = 0; - memory_space_options.inefficient_use_to_copy_ratio = 0.5; - memory_space_options.always_spill_to_default_memory = false; - memory_space_options.window_prefetch_mode = WindowPrefetchMode::kNone; - memory_space_options.expanded_scoped_alternate_memory_mode = - ExpandedScopedAlternateMemoryMode::DISABLED; - memory_space_options.post_module_scoped_alternate_memory_size_in_bytes = 0; - memory_space_options.sliced_prefetch_options.set_max_slices(4); - memory_space_options.sliced_prefetch_options.set_min_bytes(1048576); - memory_space_options.sliced_prefetch_options - .set_all_slice_time_permutations_threshold(4); - memory_space_options.memory_bound_loop_optimizer_options.set_enabled(true); - memory_space_options.memory_bound_loop_optimizer_options - .set_desired_copy_ratio(0.7); - memory_space_options.memory_bound_loop_optimizer_options - .set_allow_unsatisfied_fully_pipelined_prefetch(true); - memory_space_options.memory_bound_loop_optimizer_options - .set_min_num_iterations(6); - memory_space_options.is_allowed_in_alternate_mem_fn = - [&](const HloValue& value) { - // Check if the value belongs in the entry computation. - HloInstruction* instruction = value.instruction(); - HloComputation* computation = instruction->parent(); - bool in_entry_computation = - (computation == computation->parent()->entry_computation()); - if (in_entry_computation && - instruction->opcode() == HloOpcode::kParameter) { - return value.shape().has_layout() && - value.shape().layout().memory_space() != kDefaultMemorySpace; - } - if (instruction->opcode() == HloOpcode::kCopyStart && - value.index() == ShapeIndex({2})) { - // This is a sync flag, so don't put it in alternate memory. - return false; - } - return true; - }; - MemorySpaceAssignmentBestFitRepacker best_fit_repacker( - memory_space_options.max_size_in_bytes, - memory_space_options.alignment_in_bytes, - (memory_space_options.sliced_prefetch_options.max_slices() > - memory_space_options.sliced_prefetch_options - .all_slice_time_permutations_threshold() - ? SliceTimePermutationIterator::Ty::kPreferred - : SliceTimePermutationIterator::Ty::kAll)); - memory_space_options.repacker = &best_fit_repacker; - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - auto latency_optimized_reduce_scatter_kernel = FindInstruction( - module.get(), - "ragged_latency_optimized_reduce_scatter_out_feature_matmul.1"); - // The custom kernel should be in alternate memory. - EXPECT_EQ( - latency_optimized_reduce_scatter_kernel->shape().layout().memory_space(), - kAlternateMemorySpace); - // The operands 5 and 6 of the custom kernel should be in alternate memory. - EXPECT_EQ(latency_optimized_reduce_scatter_kernel->operand(5) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(latency_optimized_reduce_scatter_kernel->operand(6) - ->shape() - .layout() - .memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, TestColoringMultipleOperands) { - absl::string_view hlo_string = R"hlo( - HloModule NegateChain, is_scheduled=true, entry_computation_layout={(f32[2,3]{1,0}, f32[2,3]{1,0:S(2)})->f32[2,3]{1,0}} - - ENTRY %NegateChain (p0: f32[2,3], p1: f32[2,3]) -> f32[2,3] { - %p0 = f32[2,3]{1,0} parameter(0) - %p1 = f32[2,3]{1,0} parameter(1) - ROOT %add = f32[2,3]{1,0} add(f32[2,3]{1,0} %p0, f32[2,3]{1,0} %p1) - })hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* add = FindInstruction(module.get(), "add"); - HloUse add_param_1_use{add, 0, {}}; - HloUse add_param_2_use{add, 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {add_param_1_use, kAlternateMemorySpace}, - {add_param_2_use, kAlternateMemorySpace}}; - memory_space_options.max_size_in_bytes = 48; - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - - HloInstruction* add_after_msa = FindInstruction(module.get(), "add"); - EXPECT_EQ(add_after_msa->operand(0)->shape().layout().memory_space(), - kAlternateMemorySpace); - EXPECT_EQ(add_after_msa->operand(1)->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, TestColoringWithLoopOptimization) { - absl::string_view hlo_string = R"hlo( -HloModule UnrolledLoop, is_scheduled=true - -ENTRY %main { - %param.0 = (f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], - f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4], f32[2,4]) parameter(0) - - %gte.0 = f32[2,4] get-tuple-element(%param.0), index=0 - %gte.1 = f32[2,4] get-tuple-element(%param.0), index=1 - %add.common = f32[2,4] add(%gte.0, %gte.1) - %tanh.common = f32[2,4] tanh(%add.common) - - // loop idx0 - %neg.0.idx0 = f32[2,4] negate(%tanh.common) - %neg.1.idx0 = f32[2,4] negate(%neg.0.idx0) - %neg.2.idx0 = f32[2,4] negate(%neg.1.idx0) - - %gte.0.idx0 = f32[2,4] get-tuple-element(%param.0), index=2 - %add.0.idx0 = f32[2,4] add(%neg.2.idx0, %gte.0.idx0) - - %gte.1.idx0 = f32[2,4] get-tuple-element(%param.0), index=3 - %add.1.idx0 = f32[2,4] add(%add.0.idx0, %gte.1.idx0) - - %tanh.idx0 = f32[2,4] tanh(%add.1.idx0) - - // loop idx1 - %neg.0.idx1 = f32[2,4] negate(%tanh.idx0) - %neg.1.idx1 = f32[2,4] negate(%neg.0.idx1) - %neg.2.idx1 = f32[2,4] negate(%neg.1.idx1) - - %gte.0.idx1 = f32[2,4] get-tuple-element(%param.0), index=4 - %add.0.idx1 = f32[2,4] add(%neg.2.idx1, %gte.0.idx1) - - %gte.1.idx1 = f32[2,4] get-tuple-element(%param.0), index=5 - %add.1.idx1 = f32[2,4] add(%add.0.idx1, %gte.1.idx1) - - %tanh.idx1 = f32[2,4] tanh(%add.1.idx1) - - // loop idx2 - %neg.0.idx2 = f32[2,4] negate(%tanh.idx1) - %neg.1.idx2 = f32[2,4] negate(%neg.0.idx2) - %neg.2.idx2 = f32[2,4] negate(%neg.1.idx2) - - %gte.0.idx2 = f32[2,4] get-tuple-element(%param.0), index=6 - %add.0.idx2 = f32[2,4] add(%neg.2.idx2, %gte.0.idx2) - - %gte.1.idx2 = f32[2,4] get-tuple-element(%param.0), index=7 - %add.1.idx2 = f32[2,4] add(%add.0.idx2, %gte.1.idx2) - - %tanh.idx2 = f32[2,4] tanh(%add.1.idx2) - - // loop idx3 - %neg.0.idx3 = f32[2,4] negate(%tanh.idx2) - %neg.1.idx3 = f32[2,4] negate(%neg.0.idx3) - %neg.2.idx3 = f32[2,4] negate(%neg.1.idx3) - - %gte.0.idx3 = f32[2,4] get-tuple-element(%param.0), index=8 - %add.0.idx3 = f32[2,4] add(%neg.2.idx3, %gte.0.idx3) - - %gte.1.idx3 = f32[2,4] get-tuple-element(%param.0), index=9 - %add.1.idx3 = f32[2,4] add(%add.0.idx3, %gte.1.idx3) - - %tanh.idx3 = f32[2,4] tanh(%add.1.idx3) - - // loop idx4 - %neg.0.idx4 = f32[2,4] negate(%tanh.idx3) - %neg.1.idx4 = f32[2,4] negate(%neg.0.idx4) - %neg.2.idx4 = f32[2,4] negate(%neg.1.idx4) - - %gte.0.idx4 = f32[2,4] get-tuple-element(%param.0), index=10 - %add.0.idx4 = f32[2,4] add(%neg.2.idx4, %gte.0.idx4) - - %gte.1.idx4 = f32[2,4] get-tuple-element(%param.0), index=11 - %add.1.idx4 = f32[2,4] add(%add.0.idx4, %gte.1.idx4) - - %tanh.idx4 = f32[2,4] tanh(%add.1.idx4) - - // loop idx5 - %neg.0.idx5 = f32[2,4] negate(%tanh.idx4) - %neg.1.idx5 = f32[2,4] negate(%neg.0.idx5) - %neg.2.idx5 = f32[2,4] negate(%neg.1.idx5) - - %gte.0.idx5 = f32[2,4] get-tuple-element(%param.0), index=12 - %add.0.idx5 = f32[2,4] add(%neg.2.idx5, %gte.0.idx5) - - %gte.1.idx5 = f32[2,4] get-tuple-element(%param.0), index=13 - %add.1.idx5 = f32[2,4] add(%add.0.idx5, %gte.1.idx5) - - %tanh.idx5 = f32[2,4] tanh(%add.1.idx5) - - // loop idx6 - %neg.0.idx6 = f32[2,4] negate(%tanh.idx5) - %neg.1.idx6 = f32[2,4] negate(%neg.0.idx6) - %neg.2.idx6 = f32[2,4] negate(%neg.1.idx6) - - %gte.0.idx6 = f32[2,4] get-tuple-element(%param.0), index=14 - %add.0.idx6 = f32[2,4] add(%neg.2.idx6, %gte.0.idx6) - - %gte.1.idx6 = f32[2,4] get-tuple-element(%param.0), index=15 - %add.1.idx6 = f32[2,4] add(%add.0.idx6, %gte.1.idx6) - - %tanh.idx6 = f32[2,4] tanh(%add.1.idx6) - - ROOT %negate.common = f32[2,4] negate(%tanh.idx6) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.memory_bound_loop_optimizer_options.set_enabled(true); - memory_space_options.memory_bound_loop_optimizer_options - .set_desired_copy_ratio(0.1); - memory_space_options.memory_bound_loop_optimizer_options - .set_allow_unsatisfied_fully_pipelined_prefetch(true); - memory_space_options.memory_bound_loop_optimizer_options - .set_min_num_iterations(4); - HloUse add_0_idx0_use{FindInstruction(module.get(), "add.0.idx0"), 1, {}}; - HloUse add_1_idx0_use{FindInstruction(module.get(), "add.1.idx0"), 1, {}}; - HloUse add_0_idx1_use{FindInstruction(module.get(), "add.0.idx1"), 1, {}}; - HloUse add_1_idx1_use{FindInstruction(module.get(), "add.1.idx1"), 1, {}}; - HloUse add_0_idx2_use{FindInstruction(module.get(), "add.0.idx2"), 1, {}}; - HloUse add_1_idx2_use{FindInstruction(module.get(), "add.1.idx2"), 1, {}}; - HloUse add_0_idx3_use{FindInstruction(module.get(), "add.0.idx3"), 1, {}}; - HloUse add_1_idx3_use{FindInstruction(module.get(), "add.1.idx3"), 1, {}}; - HloUse add_0_idx4_use{FindInstruction(module.get(), "add.0.idx4"), 1, {}}; - HloUse add_1_idx4_use{FindInstruction(module.get(), "add.1.idx4"), 1, {}}; - HloUse add_0_idx5_use{FindInstruction(module.get(), "add.0.idx5"), 1, {}}; - HloUse add_1_idx5_use{FindInstruction(module.get(), "add.1.idx5"), 1, {}}; - HloUse add_0_idx6_use{FindInstruction(module.get(), "add.0.idx6"), 1, {}}; - HloUse add_1_idx6_use{FindInstruction(module.get(), "add.1.idx6"), 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {add_0_idx0_use, kAlternateMemorySpace}, - {add_1_idx0_use, kAlternateMemorySpace}, - {add_0_idx1_use, kAlternateMemorySpace}, - {add_1_idx1_use, kAlternateMemorySpace}, - {add_0_idx2_use, kAlternateMemorySpace}, - {add_1_idx2_use, kAlternateMemorySpace}, - {add_0_idx3_use, kAlternateMemorySpace}, - {add_1_idx3_use, kAlternateMemorySpace}, - {add_0_idx4_use, kAlternateMemorySpace}, - {add_1_idx4_use, kAlternateMemorySpace}, - {add_0_idx5_use, kAlternateMemorySpace}, - {add_1_idx5_use, kAlternateMemorySpace}, - {add_0_idx6_use, kAlternateMemorySpace}, - {add_1_idx6_use, kAlternateMemorySpace}}; - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - std::vector alternate_memory_uses = { - "add.0.idx0", "add.1.idx0", "add.0.idx1", "add.1.idx1", "add.0.idx2", - "add.1.idx2", "add.0.idx3", "add.1.idx3", "add.0.idx4", "add.1.idx4", - "add.0.idx5", "add.1.idx5", "add.0.idx6", "add.1.idx6"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/alternate_memory_uses, - /*operand_number=*/1, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, - ScopedAllocationAccountingWhenInstructionsAreRemoved) { - absl::string_view hlo_string = R"hlo( - HloModule DoIt_S64_10_0_5_1.3, is_scheduled=true - - ENTRY %DoIt_S64_10_0_5_1.3 (p0.1: (u32[10], u32[10])) -> (u32[5], u32[5]) { - %p0.1 = (u32[10]{0:T(128)}, u32[10]{0:T(128)}) parameter(0) - %get-tuple-element.1 = u32[10]{0:T(128)} get-tuple-element((u32[10]{0:T(128)}, u32[10]{0:T(128)}) %p0.1), index=1 - %bitcast.1 = u32[5]{0:T(128)} bitcast(u32[10]{0:T(128)} %get-tuple-element.1) - %get-tuple-element = u32[10]{0:T(128)} get-tuple-element((u32[10]{0:T(128)}, u32[10]{0:T(128)}) %p0.1), index=0 - %bitcast = u32[5]{0:T(128)} bitcast(u32[10]{0:T(128)} %get-tuple-element) - %tuple.1 = (u32[5]{0:T(128)}, u32[5]{0:T(128)}) tuple(u32[5]{0:T(128)} %bitcast, u32[5]{0:T(128)} %bitcast.1) - %tuple.3 = ((u32[5]{0:T(128)}, u32[5]{0:T(128)}), (u32[5]{0:T(128)}, u32[5]{0:T(128)})) tuple(%tuple.1, %tuple.1) - %get-tuple-element.4 = u32[5]{0:T(128)} get-tuple-element((u32[5]{0:T(128)}, u32[5]{0:T(128)}) %tuple.1), index=0 - %get-tuple-element.5 = (u32[5]{0:T(128)}, u32[5]{0:T(128)}) get-tuple-element(%tuple.3), index=0 - %get-tuple-element.6 = u32[5]{0:T(128)} get-tuple-element((u32[5]{0:T(128)}, u32[5]{0:T(128)}) %get-tuple-element.5), index=1 - %copy.2 = u32[5]{0:T(128)} copy(u32[5]{0:T(128)} %get-tuple-element.4) - %copy.3 = u32[5]{0:T(128)} copy(u32[5]{0:T(128)} %get-tuple-element.6) - ROOT %tuple.2 = (u32[5]{0:T(128)}, u32[5]{0:T(128)}) tuple(u32[5]{0:T(128)} %copy.2, u32[5]{0:T(128)} %copy.3) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - return 10; - }; - std::unique_ptr preset_assignments = - AssignMemorySpace(module.get(), std::move(options)); - EXPECT_THAT(preset_assignments->scoped_allocation_chunks(), - ::testing::Each(::testing::ResultOf( - "instruction->name()", - [](const auto& pair) { return pair.first->name(); }, - ::testing::AllOf(::testing::Not("bitcast"), - ::testing::Not("bitcast.1"))))); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchingNotEnoughReservedBytes) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(p5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 23; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector default_memory_uses = {"negate0", "add3", "add6", - "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/default_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kParameter); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchingLowConcurrentPrefetches) { - // The size of alternate memory allows for 4 concurrent block prefetches but - // the maximum number of concurrent block prefetches allowed is 2. The live - // ranges of p0 and p1 are long but their copy dones are early enough that - // other prefetches can start without violating the concurrent prefetch limit. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - add13 = f32[2,3]{1,0} add(p0, add12) - add14 = f32[2,3]{1,0} add(p1, add13) - ROOT add15 = f32[2,3]{1,0} add(p5, add14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 96; - memory_space_options.reserved_bytes_for_block_prefetches = 96; - memory_space_options.max_outstanding_block_prefetches = 2; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector alternate_memory_uses = { - "negate0", "add3", "add6", "add9", "add12", "add13", "add14", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/alternate_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, TestSingleBufferedBlockPrefetching) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(p5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector alternate_memory_uses = { - "negate0", "add3", "add6", "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/alternate_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchingWithAlisedUses) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - custom_call0 = f32[2,3]{1,0} custom-call(p0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate1 = f32[2,3]{1,0} negate(custom_call0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - custom_call13 = f32[2,3]{1,0} custom-call(custom_call0, add12), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate14 = f32[2,3]{1,0} negate(custom_call13) - ROOT add15 = f32[2,3]{1,0} add(p5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector prefetched_uses = {"custom_call0", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - std::vector aliased_uses = {"negate1", "custom_call13", - "negate14"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/aliased_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); - - std::vector default_memory_uses = {"add3", "add6", "add9", - "add12"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/default_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kParameter); -} - -TEST_F(MemorySpaceAssignmentTest, - TestBlockPrefetchingWithAliasingAndScopedAllocations) { - // Simple block prefetching test where p0 is block prefetched and aliases with - // custom_call2 and custom_call8 which is used at add13. This makes the - // effective aliased live range of the block prefetch p0 to be the entire - // program duration. We also have a scoped allocation for negate10 which - // reserves 24 bytes of scoped memory. The alternate memory is enough for - // prefetching p0 and allocating scoped memory for negate10. - - // If we allocate scoped memory before block prefetching, 24 bytes of scoped - // memory for negate10 will be allocated at offset 0. Block prefetch for p0 - // will also get offset 0, but we would not be able to extend it past the - // scoped allocation of negate10 when we try to extend it for the entire live - // range. - - // To make sure we don't allocate block prefetches below MaxScopedMemorySize - // we reserve memory for scoped allocations before block prefetching and - // process scoped allocations after block prefetching. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - custom_call2 = f32[2,3]{1,0} custom-call(p0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate3 = f32[2,3]{1,0} negate(p1) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[2,3]{1,0} negate(negate5) - negate7 = f32[2,3]{1,0} negate(negate6) - custom_call8 = f32[2,3]{1,0} custom-call(custom_call2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate9 = f32[2,3]{1,0} negate(negate7) - negate10 = f32[2,3]{1,0} negate(negate9) - negate11 = f32[2,3]{1,0} negate(negate10) - negate12 = f32[2,3]{1,0} negate(negate11) - ROOT add13 = f32[2,3]{1,0} add(custom_call8, negate12) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0"}); - memory_space_options.reserved_scoped_memory_fn = - [&](const HloInstruction* instruction, - const absl::flat_hash_set> - operands_in_alternate_memory, - const absl::flat_hash_set outputs_in_alternate_memory) { - if (instruction->name() == "negate10") { - return 24; - } - return 0; - }; - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector prefetched_uses = {"custom_call2"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - std::vector aliased_uses = {"custom_call8", "add13"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/aliased_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); -} - -TEST_F(MemorySpaceAssignmentTest, - TestBlockPrefetchingWithInputOutputAliasConfig) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - add15 = f32[2,3]{1,0} add(p5, negate14) - ROOT tuple = (f32[2,3]{1,0}, f32[2,3]{1,0}) tuple(add15, p3) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK(module->input_output_alias_config().SetUpAlias({1}, 3, {})); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector alternate_memory_uses = {"negate0", "add3", "add6", - "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/alternate_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - std::vector default_memory_uses = {"add9"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/default_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kParameter); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchingDoubleBuffered) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(p5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 48; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - - std::vector prefetched_uses = {"negate0", "add3", "add6", - "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchesWithMultipleUses) { - // When allowing multiple uses of block prefetches, we only prefetch once, - // making sure that the prefetches are scheduled before the first use time. - // This was a bug in the original implementation of block prefetching. This - // test case tests the fix. In this test the parameters p0 and p1 have - // multiple uses, we test that we prefetch them exactly once. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - add13 = f32[2,3]{1,0} add(p1, add12) - add14 = f32[2,3]{1,0} add(p0, add13) - ROOT add15 = f32[2,3]{1,0} add(p5, add14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 72; - memory_space_options.reserved_bytes_for_block_prefetches = 72; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2", "p3", "p4", "p5"}); - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - - std::vector alternate_memory_uses = { - "negate0", "add3", "add6", "add9", "add12", "add13", "add14", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/alternate_memory_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - HloInstruction* negate0 = FindInstruction(module.get(), "negate0"); - HloInstruction* add3 = FindInstruction(module.get(), "add3"); - HloInstruction* add13 = FindInstruction(module.get(), "add13"); - HloInstruction* add14 = FindInstruction(module.get(), "add14"); - // Check that the prefetch of p0 is reused for add14. - EXPECT_EQ(negate0->operand(0), add14->operand(0)); - // Check that the prefetch of p1 is reused for add13. - EXPECT_EQ(add3->operand(0), add13->operand(0)); -} - -TEST_F(MemorySpaceAssignmentTest, - TestBlockPrefetchingDoubleBufferedWithColoring) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(p5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 48; - memory_space_options.max_outstanding_block_prefetches = 10; - - memory_space_options.block_prefetched_positions = - GetHloPositions(module.get(), {"p5", "p4", "p3", "p2", "p1", "p0"}); - - HloInstruction* add15 = FindInstruction(module.get(), "add15"); - HloUse add15_negate14_use{add15, 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {add15_negate14_use, kAlternateMemorySpace}}; - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - - std::vector prefetched_uses = {"negate0", "add3", "add6", - "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - std::vector colored_uses = {"add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/colored_uses, - /*operand_number=*/1, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, - TestBlockPrefetchingDoubleBufferedWithColoringWithAliasAnalysis) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[2,3]{1,0} parameter(2) - p3 = f32[2,3]{1,0} parameter(3) - p4 = f32[2,3]{1,0} parameter(4) - p5 = f32[2,3]{1,0} parameter(5) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(p1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(p2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(p3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(p4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(p5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 48; - memory_space_options.max_outstanding_block_prefetches = 10; - - absl::flat_hash_set block_prefetched_positions; - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - const HloModule& hlo_module = alias_analysis->dataflow_analysis().module(); - HloComputation* entry_computation = hlo_module.entry_computation(); - for (HloInstruction* parameter_instruction : - entry_computation->parameter_instructions()) { - ShapeUtil::ForEachSubshape( - parameter_instruction->shape(), - [&](const Shape& subshape, const ShapeIndex& index) { - for (const HloBuffer* buffer : - alias_analysis->ComputeBuffersAt(parameter_instruction, index)) { - if (buffer->values().size() != 1) { - VLOG(1) << "Skipping buffer that aliases with multiple values: " - << buffer->ToString(); - continue; - } - const HloValue* value = *buffer->values().begin(); - if (value->GetUses().size() != 1) { - VLOG(1) << "Skipping buffer that has multiple uses: " - << buffer->ToString(); - continue; - } - block_prefetched_positions.insert(value->defining_position()); - } - }); - } - memory_space_options.block_prefetched_positions = block_prefetched_positions; - - HloInstruction* add15 = FindInstruction(module.get(), "add15"); - HloUse add15_negate14_use{add15, 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {add15_negate14_use, kAlternateMemorySpace}}; - - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - - std::vector prefetched_uses = {"negate0", "add3", "add6", - "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - std::vector colored_uses = {"add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/colored_uses, - /*operand_number=*/1, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(SlicedPrefetchTest, TestBlockPrefetchingAsyncSlicePrefetches) { - // Test block prefetches convert slices to async sliced prefetches. Each - // program input has two slices. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - slice0 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} - slice1 = f32[2,3]{1,0} slice(p1), slice={[0:2], [0:3]} - slice2 = f32[2,3]{1,0} slice(p2), slice={[0:2], [0:3]} - slice3 = f32[2,3]{1,0} slice(p0), slice={[2:4], [0:3]} - slice4 = f32[2,3]{1,0} slice(p1), slice={[2:4], [0:3]} - slice5 = f32[2,3]{1,0} slice(p2), slice={[2:4], [0:3]} - negate0 = f32[2,3]{1,0} negate(slice0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(slice1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(slice2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(slice3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(slice4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(slice5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kSlice; - }; - memory_space_options.block_prefetched_positions = - GetHloPositions(module.get(), {"p0", "p1", "p2", "slice0", "slice1", - "slice2", "slice3", "slice4", "slice5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector sliced_prefetch_uses = {"negate0", "add3", "add6", - "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/sliced_prefetch_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAsyncDone); - - // Check there are no additional prefetches apart from block prefetches. - const std::vector& instructions = - module->schedule().sequence(module->entry_computation()).instructions(); - int64_t num_prefetches = - absl::c_count_if(instructions, [&](const HloInstruction* instruction) { - return instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == - kAlternateMemorySpace && - (IsAsyncSliceDone(instruction) || - instruction->opcode() == HloOpcode::kCopyDone); - }); - EXPECT_EQ(num_prefetches, 6); -} - -TEST_F(MemorySpaceAssignmentTest, - TestBlockPrefetchingAsyncCustomFusionPrefetches) { - // Test block prefetches convert custom fusions to async custom fusion - // prefetches. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -custom_fusion_computation0 { - p0 = f32[4,3]{1,0} parameter(0) - ROOT slice0 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} -} - -custom_fusion_computation1 { - p0 = f32[4,3]{1,0} parameter(0) - ROOT slice1 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} -} - -custom_fusion_computation2 { - p0 = f32[4,3]{1,0} parameter(0) - ROOT slice2 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} -} - -custom_fusion_computation3 { - p0 = f32[4,3]{1,0} parameter(0) - ROOT slice3 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} -} - -custom_fusion_computation4 { - p0 = f32[4,3]{1,0} parameter(0) - ROOT slice4 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} -} - -custom_fusion_computation5 { - p0 = f32[4,3]{1,0} parameter(0) - ROOT slice5 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} -} - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - fusion0 = f32[2,3]{1,0} fusion(p0), kind=kCustom, calls=custom_fusion_computation0 - fusion1 = f32[2,3]{1,0} fusion(p1), kind=kCustom, calls=custom_fusion_computation1 - fusion2 = f32[2,3]{1,0} fusion(p2), kind=kCustom, calls=custom_fusion_computation2 - fusion3 = f32[2,3]{1,0} fusion(p0), kind=kCustom, calls=custom_fusion_computation3 - fusion4 = f32[2,3]{1,0} fusion(p1), kind=kCustom, calls=custom_fusion_computation4 - fusion5 = f32[2,3]{1,0} fusion(p2), kind=kCustom, calls=custom_fusion_computation5 - negate0 = f32[2,3]{1,0} negate(fusion0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(fusion1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(fusion2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(fusion3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(fusion4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(fusion5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.enable_sync_custom_fusion_replacement = true; - memory_space_options.custom_fusion_block_prefetch_operand_index_fn = - [](const HloInstruction* instruction) -> std::optional { - if (instruction->IsCustomFusion()) { - return 0; - } - return std::nullopt; - }; - memory_space_options.block_prefetched_positions = GetHloPositions( - module.get(), {"p0", "p1", "p2", "fusion0", "fusion1", "fusion2", - "fusion3", "fusion4", "fusion5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector custom_fusion_prefetch_uses = { - "negate0", "add3", "add6", "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), - /*instruction_names=*/custom_fusion_prefetch_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAsyncDone); - - // Check there are no additional prefetches apart from block prefetches. - const std::vector& instructions = - module->schedule().sequence(module->entry_computation()).instructions(); - int64_t num_prefetches = - absl::c_count_if(instructions, [&](const HloInstruction* instruction) { - return instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == - kAlternateMemorySpace && - (instruction->opcode() == HloOpcode::kAsyncDone && - instruction->async_wrapped_instruction()->IsCustomFusion()); - }); - EXPECT_EQ(num_prefetches, 6); -} - -TEST_F(SlicedPrefetchTest, TestBlockPrefetchingAsyncSliceWithMultipleUses) { - // Same as above, but slice 1 has multiple uses and p0 has multiple uses apart - // from the slice. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - slice0 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} - slice1 = f32[2,3]{1,0} slice(p1), slice={[0:2], [0:3]} - slice2 = f32[2,3]{1,0} slice(p2), slice={[0:2], [0:3]} - slice3 = f32[2,3]{1,0} slice(p0), slice={[2:4], [0:3]} - slice4 = f32[2,3]{1,0} slice(p1), slice={[2:4], [0:3]} - slice5 = f32[2,3]{1,0} slice(p2), slice={[2:4], [0:3]} - negate0 = f32[2,3]{1,0} negate(slice0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(slice1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(slice2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(slice3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - add11 = f32[2,3]{1,0} add(slice1, negate10) - add12 = f32[2,3]{1,0} add(slice4, add11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - add15 = f32[2,3]{1,0} add(slice5, negate14) - negate16 = f32[4,3]{1,0} negate(p0) - add17 = f32[4,3]{1,0} add(p0, negate16) - ROOT tuple = (f32[4,3]{1,0}, f32[2,3]{1,0}) tuple(add17, add15) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 96; - memory_space_options.reserved_bytes_for_block_prefetches = 96; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kSlice; - }; - - memory_space_options.block_prefetched_positions = - GetHloPositions(module.get(), {"p0", "p1", "p2", "slice0", "slice1", - "slice2", "slice3", "slice4", "slice5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector sliced_prefetch_uses = { - "negate0", "add3", "add6", "add9", "add11", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/sliced_prefetch_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAsyncDone); - - std::vector prefetched_uses = {"negate16", "add17"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - HloInstruction* add3 = FindInstruction(module.get(), "add3"); - HloInstruction* add11 = FindInstruction(module.get(), "add11"); - HloInstruction* negate16 = FindInstruction(module.get(), "negate16"); - HloInstruction* add17 = FindInstruction(module.get(), "add17"); - // Check slice1 is only prefetched once. - EXPECT_EQ(add3->operand(0), add11->operand(0)); - // Check p0 is only prefetched once. - EXPECT_EQ(negate16->operand(0), add17->operand(0)); - - // Check there are no additional prefetches apart from block prefetches. - const std::vector& instructions = - module->schedule().sequence(module->entry_computation()).instructions(); - int64_t num_prefetches = - absl::c_count_if(instructions, [&](const HloInstruction* instruction) { - return instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == - kAlternateMemorySpace && - (IsAsyncSliceDone(instruction) || - instruction->opcode() == HloOpcode::kCopyDone); - }); - EXPECT_EQ(num_prefetches, 7); -} - -TEST_F(SlicedPrefetchTest, - TestBlockPrefetchingAsyncDynamicSliceWithMultipleUses) { - // Test block prefetching with async dynamic slices. This test covers the - // scenario where: - // - a dynamic slice (dynamic_slice1) has multiple uses. - // - the original parameter being sliced (p0) has multiple uses apart from the - // dynamic slice. We expect that: - // - dynamic_slice1 is only prefetched once. - // - p0 is only prefetched once. - // - All uses of block prefetches are in the alternate memory space. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - dim0 = s32[] parameter(3) - dim1 = s32[] parameter(4) - dim2 = s32[] parameter(5) - dim3 = s32[] parameter(6) - dynamic_slice0 = f32[2,3]{1,0} dynamic-slice(p0, dim0, dim1), dynamic_slice_sizes={2,3} - dynamic_slice1 = f32[2,3]{1,0} dynamic-slice(p1, dim2, dim3), dynamic_slice_sizes={2,3} - dynamic_slice2 = f32[2,3]{1,0} dynamic-slice(p2, dim0, dim1), dynamic_slice_sizes={2,3} - dynamic_slice3 = f32[2,3]{1,0} dynamic-slice(p0, dim2, dim3), dynamic_slice_sizes={2,3} - dynamic_slice4 = f32[2,3]{1,0} dynamic-slice(p1, dim0, dim1), dynamic_slice_sizes={2,3} - dynamic_slice5 = f32[2,3]{1,0} dynamic-slice(p2, dim2, dim3), dynamic_slice_sizes={2,3} - negate0 = f32[2,3]{1,0} negate(dynamic_slice0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(dynamic_slice1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(dynamic_slice2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(dynamic_slice3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - add11 = f32[2,3]{1,0} add(dynamic_slice1, negate10) - add12 = f32[2,3]{1,0} add(dynamic_slice4, add11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - add15 = f32[2,3]{1,0} add(dynamic_slice5, negate14) - negate16 = f32[4,3]{1,0} negate(p0) - add17 = f32[4,3]{1,0} add(p0, negate16) - ROOT tuple = (f32[4,3]{1,0}, f32[2,3]{1,0}) tuple(add17, add15) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 96; - memory_space_options.reserved_bytes_for_block_prefetches = 96; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kDynamicSlice; - }; - - memory_space_options.block_prefetched_positions = GetHloPositions( - module.get(), - {"p0", "p1", "p2", "dynamic_slice0", "dynamic_slice1", "dynamic_slice2", - "dynamic_slice3", "dynamic_slice4", "dynamic_slice5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector sliced_prefetch_uses = { - "negate0", "add3", "add6", "add9", "add11", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/sliced_prefetch_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAsyncDone); - - std::vector prefetched_uses = {"negate16", "add17"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - HloInstruction* add3 = FindInstruction(module.get(), "add3"); - HloInstruction* add11 = FindInstruction(module.get(), "add11"); - HloInstruction* negate16 = FindInstruction(module.get(), "negate16"); - HloInstruction* add17 = FindInstruction(module.get(), "add17"); - // Check dynamic_slice1 is only prefetched once. - EXPECT_EQ(add3->operand(0), add11->operand(0)); - // Check p0 is only prefetched once. - EXPECT_EQ(negate16->operand(0), add17->operand(0)); - - // Check there are no additional prefetches apart from block prefetches. - const std::vector& instructions = - module->schedule().sequence(module->entry_computation()).instructions(); - int64_t num_prefetches = - absl::c_count_if(instructions, [&](const HloInstruction* instruction) { - return instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == - kAlternateMemorySpace && - ((instruction->opcode() == HloOpcode::kAsyncDone && - (instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice || - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kDynamicSlice)) || - instruction->opcode() == HloOpcode::kCopyDone); - }); - EXPECT_EQ(num_prefetches, 7); -} - -TEST_F(SlicedPrefetchTest, - TestBlockPrefetchingDoubleBufferedAsyncSlicePrefetches) { - // The block prefetching memory is twice the size of the slices. This will - // result in double buffering of the slices i.e. having 2 live async slice - // copies at a given time. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - slice0 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} - slice1 = f32[2,3]{1,0} slice(p1), slice={[0:2], [0:3]} - slice2 = f32[2,3]{1,0} slice(p2), slice={[0:2], [0:3]} - slice3 = f32[2,3]{1,0} slice(p0), slice={[2:4], [0:3]} - slice4 = f32[2,3]{1,0} slice(p1), slice={[2:4], [0:3]} - slice5 = f32[2,3]{1,0} slice(p2), slice={[2:4], [0:3]} - negate0 = f32[2,3]{1,0} negate(slice0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(slice1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(slice2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(slice3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - negate11 = f32[2,3]{1,0} negate(negate10) - add12 = f32[2,3]{1,0} add(slice4, negate11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - ROOT add15 = f32[2,3]{1,0} add(slice5, negate14) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.reserved_bytes_for_block_prefetches = 48; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kSlice; - }; - - memory_space_options.block_prefetched_positions = - GetHloPositions(module.get(), {"p0", "p1", "p2", "slice0", "slice1", - "slice2", "slice3", "slice4", "slice5"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector sliced_prefetch_uses = {"negate0", "add3", "add6", - "add9", "add12", "add15"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/sliced_prefetch_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAsyncDone); - - // Check there are no additional prefetches apart from block prefetches. - const std::vector& instructions = - module->schedule().sequence(module->entry_computation()).instructions(); - int64_t num_prefetches = - absl::c_count_if(instructions, [&](const HloInstruction* instruction) { - return instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == - kAlternateMemorySpace && - (IsAsyncSliceDone(instruction) || - instruction->opcode() == HloOpcode::kCopyDone); - }); - EXPECT_EQ(num_prefetches, 6); -} - -TEST_F(MemorySpaceAssignmentTest, TestScheduleCustomCallPrefetchesBasic) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,1,3]{2,1,0} parameter(0), sharding={replicated} - - prefetch_start_param0 = (f32[2,1,3]{2,1,0:S(1)}, s32[]{:T(128)S(2)}) custom-call(param0), custom_call_target="tpu_custom_call" - gte_param0_0 = f32[2,1,3]{2,1,0:S(1)} get-tuple-element(prefetch_start_param0), index=0 - gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(prefetch_start_param0), index=1 - prefetch_done_param0 = f32[2,1,3]{2,1,0:S(1)} custom-call(param0, gte_param0_0, gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - first_slice_prefetch_start_param0 = (f32[1,1,3]{2,1,0:S(1)}, s32[]{:T(128)S(2)}) custom-call(param0), custom_call_target="tpu_custom_call" - first_slice_gte_param0_0 = f32[1,1,3]{2,1,0:S(1)} get-tuple-element(first_slice_prefetch_start_param0), index=0 - first_slice_gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(first_slice_prefetch_start_param0), index=1 - first_slice_prefetch_done_param0 = f32[1,1,3]{2,1,0:S(1)} custom-call(param0, first_slice_gte_param0_0, first_slice_gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - second_slice_prefetch_start_param0 = (f32[1,1,3]{2,1,0:S(1)}, s32[]{:T(128)S(2)}) custom-call(param0), custom_call_target="tpu_custom_call" - second_slice_gte_param0_0 = f32[1,1,3]{2,1,0:S(1)} get-tuple-element(second_slice_prefetch_start_param0), index=0 - second_slice_gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(second_slice_prefetch_start_param0), index=1 - second_slice_prefetch_done_param0 = f32[1,1,3]{2,1,0:S(1)} custom-call(param0, second_slice_gte_param0_0, second_slice_gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - negate_param0 = f32[2,1,3]{2,1,0} negate(prefetch_done_param0) - add_param0 = f32[2,1,3]{2,1,0} add(prefetch_done_param0, negate_param0) - negate_param0_first_slice = f32[1,1,3]{2,1,0} negate(first_slice_prefetch_done_param0) - add_param0_first_slice = f32[1,1,3]{2,1,0} add(first_slice_prefetch_done_param0, negate_param0_first_slice) - negate_param0_second_slice = f32[1,1,3]{2,1,0} negate(second_slice_prefetch_done_param0) - add_param0_second_slice = f32[1,1,3]{2,1,0} add(second_slice_prefetch_done_param0, negate_param0_second_slice) - ROOT tuple = (f32[2,1,3]{2,1,0}, f32[1,1,3]{2,1,0}, f32[1,1,3]{2,1,0}) tuple(add_param0, add_param0_first_slice, add_param0_second_slice) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 200; - memory_space_options.reserved_bytes_for_block_prefetches = 200; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - std::vector custom_call_prefetch_instructions = { - {"param0", "prefetch_start_param0", "prefetch_done_param0"}, - {"param0", "first_slice_prefetch_start_param0", - "first_slice_prefetch_done_param0"}, - {"param0", "second_slice_prefetch_start_param0", - "second_slice_prefetch_done_param0"}}; - - memory_space_options.hlo_position_to_custom_call_prefetch_details = - GetCustomCallPrefetchDetailsMap(/*module=*/module.get(), - /*custom_call_prefetch_instructions=*/ - custom_call_prefetch_instructions); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector prefetch_uses = { - "add_param0", - "add_param0_first_slice", - "add_param0_second_slice", - "negate_param0", - "negate_param0_first_slice", - "negate_param0_second_slice", - }; - - // Check that all uses of custom call prefetches are in alternate memory - // space. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetch_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); -} - -TEST_F(MemorySpaceAssignmentTest, - TestScheduleCustomCallPrefetchesWithMultipleUses) { - // params p0, p1, p2 have sliced and non-sliced custom call prefetches. - // slice_done_1 and prefetch_done_0 have multiple uses. - // p0, p1 and p2 have uses apart from the custom call prefetches. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - - prefetch_start0 = (f32[4,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - prefetch_start0_gte_0 = f32[4,3]{1,0} get-tuple-element(prefetch_start0), index=0 - prefetch_start0_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(prefetch_start0), index=1 - prefetch_done_0 = f32[4,3]{1,0} custom-call(p0, prefetch_start0_gte_0, prefetch_start0_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - - slice_start_0 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - slice_start_0_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_0), index=0 - slice_start_0_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_0), index=1 - slice_done_0 = f32[2,3]{1,0} custom-call(p0, slice_start_0_gte_0, slice_start_0_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_1 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p1), custom_call_target="tpu_custom_call" - slice_start_1_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_1), index=0 - slice_start_1_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_1), index=1 - slice_done_1 = f32[2,3]{1,0} custom-call(p1, slice_start_1_gte_0, slice_start_1_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_2 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p2), custom_call_target="tpu_custom_call" - slice_start_2_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_2), index=0 - slice_start_2_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_2), index=1 - slice_done_2 = f32[2,3]{1,0} custom-call(p2, slice_start_2_gte_0, slice_start_2_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_3 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - slice_start_3_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_3), index=0 - slice_start_3_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_3), index=1 - slice_done_3 = f32[2,3]{1,0} custom-call(p0, slice_start_3_gte_0, slice_start_3_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_4 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p1), custom_call_target="tpu_custom_call" - slice_start_4_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_4), index=0 - slice_start_4_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_4), index=1 - slice_done_4 = f32[2,3]{1,0} custom-call(p1, slice_start_4_gte_0, slice_start_4_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_5 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p2), custom_call_target="tpu_custom_call" - slice_start_5_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_5), index=0 - slice_start_5_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_5), index=1 - slice_done_5 = f32[2,3]{1,0} custom-call(p2, slice_start_5_gte_0, slice_start_5_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - negate0 = f32[2,3]{1,0} negate(slice_done_0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(slice_done_1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(slice_done_2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(slice_done_3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - add11 = f32[2,3]{1,0} add(slice_done_1, negate10) - add12 = f32[2,3]{1,0} add(slice_done_4, add11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - add15 = f32[2,3]{1,0} add(slice_done_5, negate14) - negate16 = f32[4,3]{1,0} negate(prefetch_done_0) - add17 = f32[4,3]{1,0} add(prefetch_done_0, negate16) - add18 = f32[4,3]{1,0} add(p0, add17) - add19 = f32[4,3]{1,0} add(p1, add18) - add20 = f32[4,3]{1,0} add(p2, add19) - ROOT tuple = (f32[4,3]{1,0}, f32[2,3]{1,0}) tuple(add20, add15) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 96; - memory_space_options.reserved_bytes_for_block_prefetches = 96; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - std::vector custom_call_prefetch_instructions = { - {"p0", "prefetch_start0", "prefetch_done_0"}, - {"p0", "slice_start_0", "slice_done_0"}, - {"p0", "slice_start_3", "slice_done_3"}, - {"p1", "slice_start_1", "slice_done_1"}, - {"p1", "slice_start_4", "slice_done_4"}, - {"p2", "slice_start_2", "slice_done_2"}, - {"p2", "slice_start_5", "slice_done_5"}}; - - memory_space_options.hlo_position_to_custom_call_prefetch_details = - GetCustomCallPrefetchDetailsMap(/*module=*/module.get(), - /*custom_call_prefetch_instructions=*/ - custom_call_prefetch_instructions); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // Check that all uses of custom call prefetches are in alternate memory - // space. - std::vector prefetched_uses = {"negate0", "add3", "add6", - "add9", "add11", "add12", - "add15", "negate16", "add17"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); - - // Check that all direct uses of parameters are in default memory space. - std::vector direct_uses = {"add18", "add19", "add20"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/direct_uses, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kParameter); -} - -TEST_F(MemorySpaceAssignmentTest, - TestScheduleCustomCallPrefetchesWithAliasing) { - // params p0, p1, p2 have sliced and non-sliced custom call prefetches. - // slice_done_1 and prefetch_done_0 have multiple uses. - // p0, p1 and p2 have uses apart from the custom call prefetches. - // custom_call17 output aliases with prefetch_done_0, extending the live range - // of prefetch_done_0 by sharing the buffer with one additional hlo value. - // custom_call15 output aliases with slice_done_5 and custom_call21 - // output aliases with custom_call17, extending the live range of - // slice_done_5, by sharing the buffer with two additional hlo values. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - p1 = f32[4,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - - prefetch_start0 = (f32[4,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - prefetch_start0_gte_0 = f32[4,3]{1,0} get-tuple-element(prefetch_start0), index=0 - prefetch_start0_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(prefetch_start0), index=1 - prefetch_done_0 = f32[4,3]{1,0} custom-call(p0, prefetch_start0_gte_0, prefetch_start0_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - - slice_start_0 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - slice_start_0_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_0), index=0 - slice_start_0_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_0), index=1 - slice_done_0 = f32[2,3]{1,0} custom-call(p0, slice_start_0_gte_0, slice_start_0_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_1 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p1), custom_call_target="tpu_custom_call" - slice_start_1_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_1), index=0 - slice_start_1_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_1), index=1 - slice_done_1 = f32[2,3]{1,0} custom-call(p1, slice_start_1_gte_0, slice_start_1_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_2 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p2), custom_call_target="tpu_custom_call" - slice_start_2_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_2), index=0 - slice_start_2_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_2), index=1 - slice_done_2 = f32[2,3]{1,0} custom-call(p2, slice_start_2_gte_0, slice_start_2_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_3 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - slice_start_3_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_3), index=0 - slice_start_3_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_3), index=1 - slice_done_3 = f32[2,3]{1,0} custom-call(p0, slice_start_3_gte_0, slice_start_3_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_4 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p1), custom_call_target="tpu_custom_call" - slice_start_4_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_4), index=0 - slice_start_4_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_4), index=1 - slice_done_4 = f32[2,3]{1,0} custom-call(p1, slice_start_4_gte_0, slice_start_4_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - slice_start_5 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p2), custom_call_target="tpu_custom_call" - slice_start_5_gte_0 = f32[2,3]{1,0} get-tuple-element(slice_start_5), index=0 - slice_start_5_gte_1 = s32[]{:T(128)S(2)} get-tuple-element(slice_start_5), index=1 - slice_done_5 = f32[2,3]{1,0} custom-call(p2, slice_start_5_gte_0, slice_start_5_gte_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - negate0 = f32[2,3]{1,0} negate(slice_done_0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add3 = f32[2,3]{1,0} add(slice_done_1, negate2) - negate4 = f32[2,3]{1,0} negate(add3) - negate5 = f32[2,3]{1,0} negate(negate4) - add6 = f32[2,3]{1,0} add(slice_done_2, negate5) - negate7 = f32[2,3]{1,0} negate(add6) - negate8 = f32[2,3]{1,0} negate(negate7) - add9 = f32[2,3]{1,0} add(slice_done_3, negate8) - negate10 = f32[2,3]{1,0} negate(add9) - add11 = f32[2,3]{1,0} add(slice_done_1, negate10) - add12 = f32[2,3]{1,0} add(slice_done_4, add11) - negate13 = f32[2,3]{1,0} negate(add12) - negate14 = f32[2,3]{1,0} negate(negate13) - custom_call15 = f32[2,3]{1,0} custom-call(slice_done_5, negate14), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate16 = f32[4,3]{1,0} negate(prefetch_done_0) - custom_call17 = f32[4,3]{1,0} custom-call(prefetch_done_0, negate16), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - add18 = f32[4,3]{1,0} add(p0, custom_call17) - add19 = f32[4,3]{1,0} add(p1, add18) - add20 = f32[4,3]{1,0} add(p2, custom_call17) - custom_call21 = f32[4,3]{1,0} custom-call(add19, custom_call17), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - add22 = f32[4,3]{1,0} add(add20, custom_call21) - add23 = f32[2,3]{1,0} add(negate14, custom_call15) - ROOT tuple = (f32[4,3]{1,0}, f32[2,3]{1,0}) tuple(add22, add23) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 96; - memory_space_options.reserved_bytes_for_block_prefetches = 96; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - std::vector custom_call_prefetch_instructions = { - {"p0", "prefetch_start0", "prefetch_done_0"}, - {"p0", "slice_start_0", "slice_done_0"}, - {"p0", "slice_start_3", "slice_done_3"}, - {"p1", "slice_start_1", "slice_done_1"}, - {"p1", "slice_start_4", "slice_done_4"}, - {"p2", "slice_start_2", "slice_done_2"}, - {"p2", "slice_start_5", "slice_done_5"}}; - - memory_space_options.hlo_position_to_custom_call_prefetch_details = - GetCustomCallPrefetchDetailsMap(/*module=*/module.get(), - /*custom_call_prefetch_instructions=*/ - custom_call_prefetch_instructions); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // Check that all uses of custom call prefetches are in alternate memory - // space. - std::vector prefetched_uses = { - "negate0", "add3", "add6", "add9", "add11", - "add12", "custom_call15", "negate16", "custom_call17"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetched_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); - - // Check that all aliased uses of custom call prefetches are in alternate - // memory space. - std::vector aliased_uses = {"add18", "add20", "custom_call21", - "add22", "add23"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/aliased_uses, - /*operand_number=*/1, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); - - // Check that all direct uses of parameters are in default memory space. - std::vector direct_uses = {"add18", "add19", "add20"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/direct_uses, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kParameter); -} - -TEST_F(MemorySpaceAssignmentTest, TestCustomCallPrefetchSourceAliasing) { - // The block prefetched value (custom_call1) is aliased to the left and to the - // right, we test that all aliased values are in default memory space. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,1,3]{2,1,0} parameter(0), sharding={replicated} - - custom_call0 = f32[2,1,3]{2,1,0} custom-call(param0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate0 = f32[2,1,3]{2,1,0} negate(custom_call0) - negate1 = f32[2,1,3]{2,1,0} negate(negate0) - custom_call1 = f32[2,1,3]{2,1,0} custom-call(custom_call0, negate1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate2 = f32[2,1,3]{2,1,0} negate(custom_call1) - negate3 = f32[2,1,3]{2,1,0} negate(negate2) - - prefetch_start_param0 = (f32[2,1,3]{2,1,0:S(1)}, s32[]{:T(128)S(2)}) custom-call(custom_call1), custom_call_target="tpu_custom_call" - gte_param0_0 = f32[2,1,3]{2,1,0:S(1)} get-tuple-element(prefetch_start_param0), index=0 - gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(prefetch_start_param0), index=1 - prefetch_done_param0 = f32[2,1,3]{2,1,0:S(1)} custom-call(custom_call1, gte_param0_0, gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - first_slice_prefetch_start_param0 = (f32[1,1,3]{2,1,0:S(1)}, s32[]{:T(128)S(2)}) custom-call(custom_call1), custom_call_target="tpu_custom_call" - first_slice_gte_param0_0 = f32[1,1,3]{2,1,0:S(1)} get-tuple-element(first_slice_prefetch_start_param0), index=0 - first_slice_gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(first_slice_prefetch_start_param0), index=1 - first_slice_prefetch_done_param0 = f32[1,1,3]{2,1,0:S(1)} custom-call(custom_call1, first_slice_gte_param0_0, first_slice_gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - second_slice_prefetch_start_param0 = (f32[1,1,3]{2,1,0:S(1)}, s32[]{:T(128)S(2)}) custom-call(custom_call1), custom_call_target="tpu_custom_call" - second_slice_gte_param0_0 = f32[1,1,3]{2,1,0:S(1)} get-tuple-element(second_slice_prefetch_start_param0), index=0 - second_slice_gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(second_slice_prefetch_start_param0), index=1 - second_slice_prefetch_done_param0 = f32[1,1,3]{2,1,0:S(1)} custom-call(custom_call1, second_slice_gte_param0_0, second_slice_gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - negate_param0 = f32[2,1,3]{2,1,0} negate(prefetch_done_param0) - add_param0 = f32[2,1,3]{2,1,0} add(prefetch_done_param0, negate_param0) - negate_param0_first_slice = f32[1,1,3]{2,1,0} negate(first_slice_prefetch_done_param0) - add_param0_first_slice = f32[1,1,3]{2,1,0} add(first_slice_prefetch_done_param0, negate_param0_first_slice) - negate_param0_second_slice = f32[1,1,3]{2,1,0} negate(second_slice_prefetch_done_param0) - add_param0_second_slice = f32[1,1,3]{2,1,0} add(second_slice_prefetch_done_param0, negate_param0_second_slice) - - custom_call2 = f32[2,1,3]{2,1,0} custom-call(custom_call1, negate3), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate4 = f32[2,1,3]{2,1,0} negate(custom_call2) - negate5 = f32[2,1,3]{2,1,0} negate(negate4) - custom_call3 = f32[2,1,3]{2,1,0} custom-call(custom_call2, negate5), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - add0 = f32[2,1,3]{2,1,0} add(custom_call3, add_param0) - negate6 = f32[2,1,3]{2,1,0} negate(add0) - - ROOT tuple = (f32[2,1,3]{2,1,0}, f32[1,1,3]{2,1,0}, f32[1,1,3]{2,1,0}) tuple(negate6, add_param0_first_slice, add_param0_second_slice) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 1000; - memory_space_options.reserved_bytes_for_block_prefetches = 1000; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - std::vector custom_call_prefetch_instructions = { - {"custom_call1", "prefetch_start_param0", "prefetch_done_param0"}, - {"custom_call1", "first_slice_prefetch_start_param0", - "first_slice_prefetch_done_param0"}, - {"custom_call1", "second_slice_prefetch_start_param0", - "second_slice_prefetch_done_param0"}}; - - memory_space_options.hlo_position_to_custom_call_prefetch_details = - GetCustomCallPrefetchDetailsMap(/*module=*/module.get(), - /*custom_call_prefetch_instructions=*/ - custom_call_prefetch_instructions); - - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { - instruction_name_regex: "param0|custom_call0|custom_call1|custom_call2|custom_call3" - } - override_options { assign_first: true } - })pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - memory_space_options.msa_sort_order_overrides = - std::move(msa_sort_order_overrides); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - std::vector prefetch_uses = { - "add_param0", - "add_param0_first_slice", - "add_param0_second_slice", - "negate_param0", - "negate_param0_first_slice", - "negate_param0_second_slice", - }; - - // Check that all uses of custom call prefetches are in alternate memory - // space. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/prefetch_uses, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); - - std::vector default_memory_space_uses = { - "negate0", - "custom_call1", - "negate2", - "prefetch_start_param0", - "prefetch_done_param0", - "first_slice_prefetch_start_param0", - "first_slice_prefetch_done_param0", - "second_slice_prefetch_start_param0", - "second_slice_prefetch_done_param0", - "custom_call2", - "negate4", - "custom_call3", - "add0"}; - - // Check that all the uses of the hlo values that alias with the prefetched - // hlo value (custom_call1) are in default memory space even though they are - // higher in the sort order. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/default_memory_space_uses, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); -} - -TEST_F(SlicedPrefetchTest, TestMultiplePinnedAllocationsBug) { - // When block prefetching, finalize the original value if a sliced value is - // prefetched successfully and the original value is not, if not finalized it - // can cause re-allocation causing multiple pinned allocations for the same - // buffer. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[4,3]{1,0} parameter(0) - slice0 = f32[2,3]{1,0} slice(p0), slice={[0:2], [0:3]} - negate0 = f32[2,3]{1,0} negate(slice0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - add2 = f32[2,3]{1,0} add(slice0, negate2) - negate3 = f32[2,3]{1,0} negate(add2) - negate4 = f32[2,3]{1,0} negate(negate3) - negate5 = f32[2,3]{1,0} negate(negate4) - negate6 = f32[4,3]{1,0} negate(p0) - negate7 = f32[4,3]{1,0} negate(negate6) - negate8 = f32[4,3]{1,0} negate(negate7) - negate9 = f32[4,3]{1,0} negate(negate8) - add10 = f32[4,3]{1,0} add(p0, negate9) - ROOT tuple = (f32[2,3]{1,0}, f32[4,3]{1,0}) tuple(negate5, add10) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 256; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - memory_space_options.enable_sync_slice_replacement = true; - memory_space_options.is_async_slice_implemented_fn = - [](const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kSlice; - }; - memory_space_options.block_prefetched_positions = - GetHloPositions(module.get(), {"p0", "slice0"}); - const std::string text_proto = R"pb( - overrides { - hlo_position_matcher { instruction_name_regex: "p0" } - override_options { assign_first: true } - })pb"; - ASSERT_OK_AND_ASSIGN(auto msa_sort_order_overrides, - ParseTextProto(text_proto)); - memory_space_options.msa_sort_order_overrides = - std::move(msa_sort_order_overrides); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - std::vector sliced_prefetch_uses = {"negate0", "add2"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/sliced_prefetch_uses, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAsyncDone); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchingUsesInsideConditional) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -true_computation { - p2 = (f32[2,3]{1,0}) parameter(0) - gte0 = f32[2,3]{1,0} get-tuple-element(p2), index=0 - negate3 = f32[2,3]{1,0} negate(gte0) - ROOT negate4 = f32[2,3]{1,0} negate(negate3) -} - -false_computation { - p3 = (f32[2,3]{1,0}) parameter(0) - gte1 = f32[2,3]{1,0} get-tuple-element(p3), index=0 - negate5 = f32[2,3]{1,0} negate(gte1) - ROOT negate6 = f32[2,3]{1,0} negate(negate5) -} - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = pred[] parameter(1) - negate0 = f32[2,3]{1,0} negate(p0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - tuple = (f32[2,3]{1,0}) tuple(p0) - cond = f32[2,3]{1,0} conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - ROOT add0 = f32[2,3]{1,0} add(cond, negate2) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/{"negate0"}, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/{"negate3", "negate5"}, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kGetTupleElement); -} - -TEST_F(MemorySpaceAssignmentTest, - TestExplicitPrefetchingUsesInsideConditional) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -true_computation { - p2 = (f32[2,3]{1,0}) parameter(0) - gte0 = f32[2,3]{1,0} get-tuple-element(p2), index=0 - negate3 = f32[2,3]{1,0} negate(gte0) - ROOT negate4 = f32[2,3]{1,0} negate(negate3) -} - -false_computation { - p3 = (f32[2,3]{1,0}) parameter(0) - gte1 = f32[2,3]{1,0} get-tuple-element(p3), index=0 - negate5 = f32[2,3]{1,0} negate(gte1) - ROOT negate6 = f32[2,3]{1,0} negate(negate5) -} - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = pred[] parameter(1) - - prefetch_start_param0 = (f32[2,3]{1,0}, s32[]{:T(128)S(2)}) custom-call(p0), custom_call_target="tpu_custom_call" - gte_param0_0 = f32[2,3]{1,0} get-tuple-element(prefetch_start_param0), index=0 - gte_param0_1 = s32[]{:T(128)S(2)} get-tuple-element(prefetch_start_param0), index=1 - prefetch_done_param0 = f32[2,3]{1,0} custom-call(p0, gte_param0_0, gte_param0_1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (1, {})} - - negate0 = f32[2,3]{1,0} negate(prefetch_done_param0) - negate1 = f32[2,3]{1,0} negate(negate0) - negate2 = f32[2,3]{1,0} negate(negate1) - tuple = (f32[2,3]{1,0}) tuple(prefetch_done_param0) - cond = f32[2,3]{1,0} conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - ROOT add0 = f32[2,3]{1,0} add(cond, negate2) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - std::vector custom_call_prefetch_instructions = { - {"p0", "prefetch_start_param0", "prefetch_done_param0"}}; - - memory_space_options.hlo_position_to_custom_call_prefetch_details = - GetCustomCallPrefetchDetailsMap(/*module=*/module.get(), - /*custom_call_prefetch_instructions=*/ - custom_call_prefetch_instructions); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/{"negate0"}, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCustomCall); - - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/{"negate3", "negate5"}, - /*operand_number=*/0, - /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kGetTupleElement); -} - -TEST_F(MemorySpaceAssignmentTest, NoPrefetchWithBandwidthLimitingAsyncStart) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. The prefetch will not happen because the bandwidth - // adjustment factor is 0.0. - const absl::string_view hlo_string = R"hlo( -HloModule NegateChain, is_scheduled=true, entry_computation_layout={(f32[2,3]{1,0}, f32[2,3]{1,0})->f32[2,3]{1,0}} - -ENTRY %NegateChain (p0: f32[2,3], p1: f32[2,3]) -> f32[2,3] { - %p0 = f32[2,3]{1,0} parameter(0) - %p1 = f32[2,3]{1,0} parameter(1) - %negate = f32[2,3]{1,0} negate(%p0) - %negate.1 = f32[2,3]{1,0} negate(%negate) - %negate.2 = f32[2,3]{1,0} negate(%negate.1) - %negate.3 = f32[2,3]{1,0} negate(%negate.2) - %negate.4 = f32[2,3]{1,0} negate(%negate.3) - %negate.5 = f32[2,3]{1,0} negate(%negate.4) - %negate.6 = f32[2,3]{1,0} negate(%negate.5) - ROOT %add = f32[2,3]{1,0} add(%negate.6, %p1) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - auto async_instruction_bw_adjustment_factor_fn = - [&](const HloInstruction* inst) -> std::optional { - if (inst->name() == "p1") { - return 0.0f; - } - return std::nullopt; - }; - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.async_instruction_bw_adjustment_factor_fn = - async_instruction_bw_adjustment_factor_fn; - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - HloInstruction* add = FindInstruction(module.get(), "add"); - EXPECT_EQ(add->operand(1)->shape().layout().memory_space(), - kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - ReserveColoredBuffersBeforeCrossProgramPrefetchBug) { - // Reserving alternate memory space for colored buffers first prevents any - // successful cross program prefetches. - absl::string_view hlo_string = R"hlo( - HloModule cross_program_prefetch, is_scheduled=true - - ENTRY cross_program_prefetch { - p0 = (f32[8,8]{1,0}, f32[8,2]{1,0}) parameter(0) - get-tuple-element.0 = f32[8,8]{1,0} get-tuple-element(p0), index=0 - add.0 = f32[8,8]{1,0} add(get-tuple-element.0, get-tuple-element.0) - get-tuple-element.1 = f32[8,2]{1,0} get-tuple-element(p0), index=1 - dot.0 = f32[8,2]{1,0} dot(add.0, get-tuple-element.1), lhs_contracting_dims={1}, rhs_contracting_dims={0} - negate.1 = f32[8,2]{1,0} negate(dot.0) - negate.2 = f32[8,2]{1,0} negate(negate.1) - ROOT dot.1 = f32[2,2]{1,0} dot(negate.2, get-tuple-element.1), lhs_contracting_dims={0}, rhs_contracting_dims={0} - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - Options options = DefaultMemorySpaceOptions(); - options.enable_cross_program_prefetch = true; - options.max_size_in_bytes = 256; - - HloInstruction* dot0 = FindInstruction(module.get(), "dot.0"); - HloUse dot0_use_add0 = {dot0, 0, {}}; - options.allocate_colored_buffers_early = false; - options.buffer_colorings = {{dot0_use_add0, kAlternateMemorySpace}}; - - AssignMemorySpace(module.get(), std::move(options)); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"dot.0"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kAdd); - auto cross_program_prefetches = module->CrossProgramPrefetches(); - ASSERT_EQ(cross_program_prefetches.size(), 0); -} - -TEST_F(MemorySpaceAssignmentTest, PrefetchWithoutBandwidthLimitingAsyncStart) { - // The negate chain is long enough for asynchronous copy to be inserted - // between p1 and add. The prefetch will happen because the bandwidth - // adjustment factor is 1.0. - const absl::string_view hlo_string = R"hlo( -HloModule NegateChain, is_scheduled=true, entry_computation_layout={(f32[2,3]{1,0}, f32[2,3]{1,0})->f32[2,3]{1,0}} - -ENTRY %NegateChain (p0: f32[2,3], p1: f32[2,3]) -> f32[2,3] { - %p0 = f32[2,3]{1,0} parameter(0) - %p1 = f32[2,3]{1,0} parameter(1) - %negate = f32[2,3]{1,0} negate(%p0) - %negate.1 = f32[2,3]{1,0} negate(%negate) - %negate.2 = f32[2,3]{1,0} negate(%negate.1) - %negate.3 = f32[2,3]{1,0} negate(%negate.2) - %negate.4 = f32[2,3]{1,0} negate(%negate.3) - %negate.5 = f32[2,3]{1,0} negate(%negate.4) - %negate.6 = f32[2,3]{1,0} negate(%negate.5) - ROOT %add = f32[2,3]{1,0} add(%negate.6, %p1) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - auto async_instruction_bw_adjustment_factor_fn = - [&](const HloInstruction* inst) -> std::optional { - if (inst->name() == "p1") { - return 1.0f; - } - return std::nullopt; - }; - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.async_instruction_bw_adjustment_factor_fn = - async_instruction_bw_adjustment_factor_fn; - XLA_VLOG_LINES(3, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(3, "After MSA: \n" + module->ToString()); - HloInstruction* add = FindInstruction(module.get(), "add"); - EXPECT_EQ(add->operand(1)->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, TestForceImmediateEvictionBug) { - // * negate0 output is colored in alternate memory space - // * negate1 output is colored in alternate memory space - // * negate2 is pre-colored in alternate memory space - // Since the size of alternate memory space is 48 bytes, which is enough for - // exactly one value, this forces immediate eviction of both neagte0 and - // negate1. The incorrect copy start time of eviction caused a scheduling bug - // that resulted evictions for negate0 and negate1 to be overlapped. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - negate0 = f32[3,4]{1,0} negate(p0) - negate1 = f32[3,4]{1,0} negate(p1) - negate2 = f32[3,4]{1,0:S(1)} negate(negate0) - negate3 = f32[3,4]{1,0} negate(negate1) - ROOT add0 = f32[3,4]{1,0} add(negate2, negate3) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* negate0 = FindInstruction(module.get(), "negate0"); - HloInstruction* negate1 = FindInstruction(module.get(), "negate1"); - HloPosition negate0_position{negate0, {}}; - HloPosition negate1_position{negate1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate0_position, kAlternateMemorySpace}, - {negate1_position, kAlternateMemorySpace}}; - memory_space_options.max_size_in_bytes = 48; - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - std::vector instruction_names = {"negate0", "negate1", - "negate2"}; - CheckMemorySpaceForInstructionNames(module.get(), instruction_names, - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, TestSchedulingImmediateCopiesBug) { - // * negate3 output is colored in alternate memory space - // * add0 use of negate2 is colored in alternate memory space - // * alternate memory size is 48 bytes, enough for exactly one value - // * this forces immediate eviction of negate3 and just in time prefetch of - // negate2. - // * msa should schedule the eviction start and done of negate3 first followed - // by the immediate prefetch for negate2. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - negate0 = f32[3,4]{1,0} negate(p0) - negate1 = f32[3,4]{1,0} negate(p1) - negate2 = f32[3,4]{1,0} negate(negate0) - negate3 = f32[3,4]{1,0} negate(negate1) - add0 = f32[3,4]{1,0} add(negate2, negate3) - negate4 = f32[3,4]{1,0} negate(add0) - negate5 = f32[3,4]{1,0} negate(negate4) - ROOT negate6 = f32[3,4]{1,0} negate(negate5) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - - HloInstruction* add0 = FindInstruction(module.get(), "add0"); - HloInstruction* negate3 = FindInstruction(module.get(), "negate3"); - HloPosition negate3_position{negate3, {}}; - HloUse add0_use_of_negate2{add0, 0, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate3_position, kAlternateMemorySpace}, - {add0_use_of_negate2, kAlternateMemorySpace}}; - memory_space_options.max_size_in_bytes = 48; - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - std::vector output_colored = {"negate3"}; - CheckMemorySpaceForInstructionNames(module.get(), output_colored, - kAlternateMemorySpace); - std::vector input_colored = {"add0"}; - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), input_colored, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, ScheduleDonesBeforeStarts) { - // Alternate memory size is 48 bytes, enough for exactly one value. - // Eviction done of negate3 should be scheduled before prefetch start of - // negate2, to free up the alternate memory space for prefetch of negate2. - - // * negate2 is colored in alternate memory space till the end of sub0. - // * negate3 is colored in default memory space till the end of sub0. - // * eviction done of negate2 should be scheduled right before mul0. - // * prefetch start of negate2 should be scheduled right after sub0. - // * they should be scheduled in the order of: negate3 eviction done, - // negate2 prefetch start. - - // NOTE: In MSA we conservatively schedule the eviction done one instruction - // earlier, which is why this corner case has never been triggered. This only - // commes into effect for immediate evictions. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - negate0 = f32[3,4]{1,0} negate(p0) - negate1 = f32[3,4]{1,0} negate(p1) - negate2 = f32[3,4]{1,0} negate(negate0) // output color negate2 in alternate memory space - negate3 = f32[3,4]{1,0} negate(negate1) // output color negate3 in default memory space - tanh0 = f32[3,4]{1,0} tanh(negate0) - add0 = f32[3,4]{1,0} add(negate2, negate3) // input color negate2 in alternate memory space and negate3 in default memory space - sub0 = f32[3,4]{1,0} subtract(negate2, negate3) // input color negate2 in alternate memory space and negate3 in default memory space - mul0 = f32[3,4]{1,0} multiply(negate2, negate3) // input color negate2 in default memory space and negate3 in alternate memory space - div0 = f32[3,4]{1,0} divide(negate2, negate3) // input color negate2 in default memory space and negate3 in alternate memory space - add1 = f32[3,4]{1,0} add(add0, sub0) - add2 = f32[3,4]{1,0} add(mul0, div0) - add3 = f32[3,4]{1,0} add(add1, add2) - ROOT add4 = f32[3,4]{1,0} add(tanh0, add3) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - - HloPosition negate2_position = {FindInstruction(module.get(), "negate2"), {}}; - HloPosition negate3_position = {FindInstruction(module.get(), "negate3"), {}}; - HloUse add0_use_of_negate2{FindInstruction(module.get(), "add0"), 0, {}}; - HloUse sub0_use_of_negate2{FindInstruction(module.get(), "sub0"), 0, {}}; - HloUse mul0_use_of_negate2{FindInstruction(module.get(), "mul0"), 0, {}}; - HloUse div0_use_of_negate2{FindInstruction(module.get(), "div0"), 0, {}}; - HloUse add0_use_of_negate3{FindInstruction(module.get(), "add0"), 1, {}}; - HloUse sub0_use_of_negate3{FindInstruction(module.get(), "sub0"), 1, {}}; - HloUse mul0_use_of_negate3{FindInstruction(module.get(), "mul0"), 1, {}}; - HloUse div0_use_of_negate3{FindInstruction(module.get(), "div0"), 1, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate2_position, kAlternateMemorySpace}, - {add0_use_of_negate2, kAlternateMemorySpace}, - {sub0_use_of_negate2, kAlternateMemorySpace}, - {mul0_use_of_negate2, kDefaultMemorySpace}, - {div0_use_of_negate2, kDefaultMemorySpace}, - {negate3_position, kDefaultMemorySpace}, - {add0_use_of_negate3, kDefaultMemorySpace}, - {sub0_use_of_negate3, kDefaultMemorySpace}, - {mul0_use_of_negate3, kAlternateMemorySpace}, - {div0_use_of_negate3, kAlternateMemorySpace}}; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "negate2") { - priority = 1; - } else if (x.buffer->instruction()->name() == "negate3") { - priority = 2; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 100); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - CheckMemorySpaceForInstructionNames(module.get(), {"negate2"}, - kAlternateMemorySpace); - CheckMemorySpaceForInstructionNames(module.get(), {"negate3"}, - kDefaultMemorySpace); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"add0", "sub0"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kNegate); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"add0", "sub0"}, - /*operand_number=*/1, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kNegate); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"mul0", "div0"}, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"mul0", "div0"}, - /*operand_number=*/1, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, TestColoringSyncReplacementCandidateUses) { - // * negate0 is colored in alternate memory space - // * negate2 and tanh0 uses of negate0 are colored in alternate memory space - // * add2 and add4 uses of copy0 are colored in alternate memory space - // * As a result, we should see an eviction of negate0 in default memory - // space because the alternate memory space is not big enough to hold - // negate0 and the copy of negate0 at the same time. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - negate0 = f32[3,4]{1,0} negate(p0) - negate1 = f32[3,4]{1,0} negate(p1) - negate2 = f32[3,4]{1,0} negate(negate0) - negate3 = f32[3,4]{1,0} negate(negate1) - tanh0 = f32[3,4]{1,0} tanh(negate0) - add0 = f32[3,4]{1,0} add(negate2, negate3) - copy0 = f32[3,4]{1,0} copy(negate0) - add1 = f32[3,4]{1,0} add(tanh0, add0) - add2 = f32[3,4]{1,0} add(copy0, add1) - add3 = f32[3,4]{1,0} add(negate0, add2) - add4 = f32[3,4]{1,0} add(copy0, add3) - ROOT negate4 = f32[3,4]{1,0} negate(add4) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.enable_sync_copy_replacement = true; - memory_space_options.enable_sync_slice_replacement = true; - HloUse add2_use_of_copy0{FindInstruction(module.get(), "add2"), 0, {}}; - HloUse add4_use_of_copy0{FindInstruction(module.get(), "add4"), 0, {}}; - HloUse negate2_use_of_negate0{ - FindInstruction(module.get(), "negate2"), 0, {}}; - HloUse tanh0_use_of_negate0{FindInstruction(module.get(), "tanh0"), 0, {}}; - HloPosition negate0_position{FindInstruction(module.get(), "negate0"), {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate0_position, kAlternateMemorySpace}, - {negate2_use_of_negate0, kAlternateMemorySpace}, - {tanh0_use_of_negate0, kAlternateMemorySpace}, - {add2_use_of_copy0, kAlternateMemorySpace}, - {add4_use_of_copy0, kAlternateMemorySpace}}; - - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "negate0") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 100); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // Check that negate0 is alternate memory space - CheckMemorySpaceForInstructionNames(module.get(), {"negate0"}, - kAlternateMemorySpace); - - // Check that negate0's uses at negate2 and tanh0 are from the pinned - // allocation in alternate memory space. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"negate2", "tanh0"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kNegate); - - // Check that copy0 is converted to async copy to alternate memory space. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"add2", "add4"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); - - // Check that negate0's use at add3 is from an eviction in default - // memory space. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"add3"}, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kCopyDone); -} - -TEST_F(MemorySpaceAssignmentTest, - TestColoringSyncSliceReplacementCandidateUses) { - // * negate0 is colored in alternate memory space - // * negate2 and tanh0 uses of negate0 are colored in alternate memory space - // * negate4 and tanh1 uses of slice0 are colored in alternate memory space - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[3,4]{1,0} parameter(0) - p1 = f32[3,4]{1,0} parameter(1) - negate0 = f32[3,4]{1,0} negate(p0) - negate1 = f32[3,4]{1,0} negate(p1) - negate2 = f32[3,4]{1,0} negate(negate0) - negate3 = f32[3,4]{1,0} negate(negate1) - tanh0 = f32[3,4]{1,0} tanh(negate0) - add0 = f32[3,4]{1,0} add(negate2, negate3) - slice0 = f32[3,2]{1,0} slice(negate0), slice={[0:3], [0:2]} - add1 = f32[3,4]{1,0} add(tanh0, add0) - negate4 = f32[3,2]{1,0} negate(slice0) - add3 = f32[3,4]{1,0} add(negate0, add1) - tanh1 = f32[3,2]{1,0} tanh(slice0) - add5 = f32[3,2]{1,0} add(negate4, tanh1) - ROOT tuple0 = (f32[3,4], f32[3,2]) tuple(add3, add5) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.enable_sync_copy_replacement = true; - memory_space_options.enable_sync_slice_replacement = true; - HloUse negate4_use_of_slice0{FindInstruction(module.get(), "negate4"), 0, {}}; - HloUse tanh1_use_of_slice0{FindInstruction(module.get(), "tanh1"), 0, {}}; - HloUse negate2_use_of_negate0{ - FindInstruction(module.get(), "negate2"), 0, {}}; - HloUse tanh0_use_of_negate0{FindInstruction(module.get(), "tanh0"), 0, {}}; - HloPosition negate0_position{FindInstruction(module.get(), "negate0"), {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate0_position, kAlternateMemorySpace}, - {negate2_use_of_negate0, kAlternateMemorySpace}, - {tanh0_use_of_negate0, kAlternateMemorySpace}, - {negate4_use_of_slice0, kAlternateMemorySpace}, - {tanh1_use_of_slice0, kAlternateMemorySpace}}; - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // Check that negate0 is alternate memory space - CheckMemorySpaceForInstructionNames(module.get(), {"negate0"}, - kAlternateMemorySpace); - // Check that slice0 uses at negate4 and tanh1 and negate0 uses at negate2 and - // tanh0 read from alternate memory space. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"negate2", "tanh0", "negate4", "tanh1"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, TestReresrvingPendingAllocations) { - // * negate0, negate3, and negate4 require contiguous allocations. - // * negate0 and negate4 are colored in alternate memory space. - // * Sort order is negate3, negate4, negate0 followed by other buffers. - // * When we force fail the allocation of negate0, the reserved allocation - // should be re-reserved during repacking. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,3]{1,0} parameter(1), sharding={replicated} - param1 = f32[4,3]{1,0} parameter(0), sharding={replicated} - param2 = f32[2,3]{1,0} parameter(2), sharding={replicated} - negate0 = f32[4,3]{1,0} negate(param1) - negate1 = f32[4,3]{1,0} negate(negate0) - negate2 = f32[4,3]{1,0} negate(negate1) - negate3 = f32[2,3]{1,0} negate(param0) - add0 = f32[4,3]{1,0} add(negate0, negate2) - negate4 = f32[2,3]{1,0} negate(param2) - tanh0 = f32[4,3]{1,0} tanh(add0) - negate5 = f32[2,3]{1,0} negate(negate4) - add1 = f32[2,3]{1,0} add(negate3,negate5) - ROOT tuple0 = (f32[2,3]{1,0}, f32[4,3]{1,0}) tuple(add1, tanh0) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 72; - memory_space_options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - return position.instruction->name() == "negate0" || - position.instruction->name() == "negate3" || - position.instruction->name() == "negate4"; - }; - HloPosition negate0_position{FindInstruction(module.get(), "negate0"), {}}; - HloPosition negate4_position{FindInstruction(module.get(), "negate4"), {}}; - memory_space_options.max_retries = 2; - memory_space_options.max_repacks = 2; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate0_position, kAlternateMemorySpace}, - {negate4_position, kAlternateMemorySpace}}; - absl::flat_hash_map, int64_t> repack_map; - auto check_fun = [](absl::Span allocations) { - // 3 allocations: 2 committed allocations for negate3 and negate4 - // 1 re-reserved allocation for negate0 - EXPECT_EQ(allocations.size(), 3); - }; - FakeMemorySpaceAssignmentRepacker repacker = - FakeMemorySpaceAssignmentRepacker(repack_map, check_fun, true); - - memory_space_options.repacker = &repacker; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "negate3") { - priority = 1; - } - if (x.buffer->instruction()->name() == "negate4") { - priority = 2; - } - if (x.buffer->instruction()->name() == "negate0") { - priority = 3; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - return lookup(lhs) < lookup(rhs); - }; - - memory_space_options.allocation_result_modifier_testing_fn = - [](const AllocationRequest& request, AllocationResult& result, - int64_t retry_number) { - if (request.allocation_value->position().instruction->name() == - "negate0" && - retry_number <= 0 && result == AllocationResult::kSuccess) { - result = AllocationResult::kFailOutOfMemory; - } - }; - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 100); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - CheckMemorySpaceForInstructionNames( - module.get(), {"negate0", "negate4", "negate3"}, kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - TestColoringWithContiguouslyAllocatedBuffers) { - // * All the custom calls alias with each other. The buffer live ranges starts - // at custom_call0 and ends at the last use of custom_call3 at negate8. - // * All the custom call buffers require contiguous allocation. - // * Coloring the buffer at one instruction site at negate5 (use of - // custom_call2) puts the entire live range of the buffer in alternate - // memory space. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,1,3]{2,1,0} parameter(0), sharding={replicated} - negate0 = f32[2,1,3]{2,1,0} negate(param0) - custom_call0 = f32[2,1,3]{2,1,0} custom-call(negate0), custom_call_target="tpu_custom_call" - negate1 = f32[2,1,3]{2,1,0} negate(custom_call0) - negate2 = f32[2,1,3]{2,1,0} negate(negate1) - custom_call1 = f32[2,1,3]{2,1,0} custom-call(custom_call0, negate2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate3 = f32[2,1,3]{2,1,0} negate(custom_call1) - negate4 = f32[2,1,3]{2,1,0} negate(negate3) - custom_call2 = f32[2,1,3]{2,1,0} custom-call(custom_call1, negate4), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate5 = f32[2,1,3]{2,1,0} negate(custom_call2) - negate6 = f32[2,1,3]{2,1,0} negate(negate5) - custom_call3 = f32[2,1,3]{2,1,0} custom-call(custom_call2, negate6), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate7 = f32[2,1,3]{2,1,0} negate(negate6) - negate8 = f32[2,1,3]{2,1,0} negate(custom_call3) - ROOT add0 = f32[2,1,3]{2,1,0} add(negate7, negate8) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - return position.instruction->opcode() == HloOpcode::kCustomCall; - }; - HloUse negate5_use_of_custom_call2{ - FindInstruction(module.get(), "negate5"), 0, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate5_use_of_custom_call2, kAlternateMemorySpace}, - }; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "negate4") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 100); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - CheckMemorySpaceForInstructionNames( - module.get(), - {"custom_call0", "custom_call1", "custom_call2", "custom_call3"}, - kAlternateMemorySpace); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), - {"negate1", "negate3", "negate5", "negate8", "custom_call1", - "custom_call2", "custom_call3"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - HloOpcode::kCustomCall); -} - -TEST_F(MemorySpaceAssignmentTest, TestColoringWithContiguousBuffersExtended) { - // * All the custom calls alias with each other and with negate0. The buffer - // live ranges starts at negate0 and ends at the last use of custom_call3 at - // negate8. - // * Only the custom call buffers require contiguous allocation (negate0 does - // not require contiguous allocation). - // * Coloring the buffer at one instruction site at negate5 (use of - // custom_call2) puts the entire live range of the buffer in alternate - // memory space. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - param0 = f32[2,1,3]{2,1,0} parameter(0), sharding={replicated} - negate0 = f32[2,1,3]{2,1,0} negate(param0) - custom_call0 = f32[2,1,3]{2,1,0} custom-call(negate0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate1 = f32[2,1,3]{2,1,0} negate(custom_call0) - negate2 = f32[2,1,3]{2,1,0} negate(negate1) - custom_call1 = f32[2,1,3]{2,1,0} custom-call(custom_call0, negate2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate3 = f32[2,1,3]{2,1,0} negate(custom_call1) - negate4 = f32[2,1,3]{2,1,0} negate(negate3) - custom_call2 = f32[2,1,3]{2,1,0} custom-call(custom_call1, negate4), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate5 = f32[2,1,3]{2,1,0} negate(custom_call2) - negate6 = f32[2,1,3]{2,1,0} negate(negate5) - custom_call3 = f32[2,1,3]{2,1,0} custom-call(custom_call2, negate6), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate7 = f32[2,1,3]{2,1,0} negate(negate6) - negate8 = f32[2,1,3]{2,1,0} negate(custom_call3) - ROOT add0 = f32[2,1,3]{2,1,0} add(negate7, negate8) -})hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 48; - memory_space_options.position_requires_contiguous_allocation_fn = - [](const HloPosition& position) { - return position.instruction->opcode() == HloOpcode::kCustomCall; - }; - HloUse negate5_use_of_custom_call2{ - FindInstruction(module.get(), "negate5"), 0, {}}; - memory_space_options.allocate_colored_buffers_early = false; - memory_space_options.buffer_colorings = { - {negate5_use_of_custom_call2, kAlternateMemorySpace}, - }; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "negate4") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 100); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - CheckMemorySpaceForInstructionNames( - module.get(), - {"negate0", "custom_call0", "custom_call1", "custom_call2", - "custom_call3"}, - kAlternateMemorySpace); - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), - {"negate1", "negate3", "negate5", "negate8", "custom_call0", - "custom_call1", "custom_call2", "custom_call3"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - ConditionalAliasedInputOutputAllocatedInAlternateMemory) { - // Conditional outputs and inputs alias with each other, and the outputs are - // first in the sort order. They should all get alternate memory allocations. - // Alternate memory size is 40 bytes, just enough for the both conditional - // outputs. - // This is a special case where one input to each branch computation aliases - // with one of the outputs in the branch and other output for the branch is a - // new buffer. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - computation0 { - param0 = (f32[3,2]{1,0}) parameter(0) - gte0 = f32[3,2]{1,0} get-tuple-element(param0), index=0 - cc0 = f32[3,2]{1,0} custom-call(gte0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - cc1 = f32[3,2]{1,0} custom-call(cc0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - empty0 = f32[2,2]{1,0} custom-call(), custom_call_target="AllocateBuffer" - ROOT tuple2 = (f32[3,2], f32[2,2]) tuple(cc1, empty0) - } - - computation1 { - param1 = (f32[2,2]{1,0}) parameter(0) - get1 = f32[2,2]{1,0} get-tuple-element(param1), index=0 - cc2 = f32[2,2]{1,0} custom-call(get1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - cc3 = f32[2,2]{1,0} custom-call(cc2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - empty1 = f32[3,2]{1,0} custom-call(), custom_call_target="AllocateBuffer" - ROOT tuple3 = (f32[3,2], f32[2,2]) tuple(empty1, cc3) - } - - ENTRY entry { - p0 = f32[3,2] parameter(0) - p1 = f32[2,2] parameter(1) - p2 = pred[] parameter(2) - negate0 = f32[3,2]{1,0} negate(p0) - negate1 = f32[2,2]{1,0} negate(p1) - tuple0 = (f32[3,2]{1,0}) tuple(negate0) - tuple1 = (f32[2,2]{1,0}) tuple(negate1) - conditional = (f32[3,2], f32[2,2]) conditional(p2, tuple0, tuple1), true_computation=computation0, false_computation=computation1 - gte2 = f32[3,2]{1,0} get-tuple-element(conditional), index=0 - gte3 = f32[2,2]{1,0} get-tuple-element(conditional), index=1 - negate6 = f32[3,2]{1,0} negate(gte2) - negate7 = f32[2,2]{1,0} negate(gte3) - ROOT tuple2 = (f32[3,2]{1,0}, f32[2,2]{1,0}) tuple(negate6, negate7) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 40; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - absl::flat_hash_set priority_instructions = { - "empty0", "empty1", "cc1", "cc3"}; - if (priority_instructions.contains(x.buffer->instruction()->name())) { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // Verify conditional outputs and aliasing inputs are in alternate memory. - CheckMemorySpaceForInstructionNames( - module.get(), {"negate0", "negate1", "empty0", "empty1", "cc1", "cc3"}, - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, AllocateConditionalOutputsInAlternateMemory) { - // Conditional outputs will be allocated first. There is just enough space in - // alternate memory to fit them. - // custom-call outputs will also get alternate memory allocations, because - // they alias with the conditional outputs. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - computation0 { - branch0_param0 = (f32[3,2]{1,0}) parameter(0) - branch0_gte0 = f32[3,2]{1,0} get-tuple-element(branch0_param0), index=0 - branch0_negate0 = f32[3,2]{1,0} negate(branch0_gte0) - branch0_negate1 = f32[3,2]{1,0} negate(branch0_negate0) - branch0_empty0 = f32[2,2]{1,0} custom-call(), custom_call_target="AllocateBuffer" - ROOT tuple2 = (f32[3,2], f32[2,2]) tuple(branch0_negate1, branch0_empty0) - } - - computation1 { - branch1_param0 = (f32[2,2]{1,0}) parameter(0) - branch1_gte0 = f32[2,2]{1,0} get-tuple-element(branch1_param0), index=0 - branch1_negate0 = f32[2,2]{1,0} negate(branch1_gte0) - branch1_negate1 = f32[2,2]{1,0} negate(branch1_negate0) - branch1_empty0 = f32[3,2]{1,0} custom-call(), custom_call_target="AllocateBuffer" - ROOT tuple3 = (f32[3,2], f32[2,2]) tuple(branch1_empty0, branch1_negate1) - } - - ENTRY entry { - p0 = f32[3,2] parameter(0) - p1 = f32[2,2] parameter(1) - p2 = pred[] parameter(2) - negate0 = f32[3,2]{1,0} negate(p0) - negate1 = f32[2,2]{1,0} negate(p1) - tuple0 = (f32[3,2]{1,0}) tuple(negate0) - tuple1 = (f32[2,2]{1,0}) tuple(negate1) - conditional = (f32[3,2], f32[2,2]) conditional(p2, tuple0, tuple1), true_computation=computation0, false_computation=computation1 - gte2 = f32[3,2]{1,0} get-tuple-element(conditional), index=0 - gte3 = f32[2,2]{1,0} get-tuple-element(conditional), index=1 - custom_call6 = f32[3,2]{1,0} custom-call(gte2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - custom_call7 = f32[2,2]{1,0} custom-call(gte3), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - negate8 = f32[3,2]{1,0} negate(custom_call6) - negate9 = f32[2,2]{1,0} negate(custom_call7) - ROOT tuple4 = (f32[3,2]{1,0}, f32[2,2]{1,0}) tuple(negate8, negate9) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 40; - memory_space_options.max_outstanding_prefetches = 0; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - absl::flat_hash_set priority_instructions = { - "branch0_empty0", "branch1_empty0", "branch0_negate1", - "branch1_negate1"}; - if (priority_instructions.contains(x.buffer->instruction()->name())) { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // verify conditional outputs and aliasing custom call outputs are in - // alternate memory. - CheckMemorySpaceForInstructionNames( - module.get(), - {"branch0_empty0", "branch1_empty0", "branch0_negate1", "branch1_negate1", - "custom_call6", "custom_call7"}, - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalWithoutTupleAllocationTest) { - // Conditional branch computation outputs are not put in a tuple, they won't - // be allocated in alternate memory even though they are first in sort order. - // This is because they are produced at the root instruction of the - // computation and thus their allocation values have no uses. - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -true_computation { - branch0_param0 = (f32[2,3]) parameter(0) - gte1 = f32[2,3] get-tuple-element(branch0_param0), index=0 - neg0 = f32[2,3] negate(gte1) - ROOT neg1 = f32[2,3] negate(neg0) -} - -false_computation { - branch1_param0 = (f32[2,3]) parameter(0) - gte2 = f32[2,3] get-tuple-element(branch1_param0), index=0 - neg2 = f32[2,3] negate(gte2) - ROOT neg3 = f32[2,3] negate(neg2) -} - -ENTRY entry { - p0 = f32[2,3] parameter(0) - p1 = pred[] parameter(1) - tuple = (f32[2,3]) tuple(p0) - conditional = f32[2,3] conditional(p1, tuple, tuple), true_computation=true_computation, false_computation=false_computation - neg4 = f32[2,3] negate(conditional) - ROOT neg5 = f32[2,3] negate(neg4) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 1000; - memory_space_options.max_outstanding_prefetches = 0; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "neg1" || - x.buffer->instruction()->name() == "neg3") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // verify conditional outputs are not in alternate memory. - CheckMemorySpaceForInstructionNames(module.get(), {"neg1", "neg3"}, - kDefaultMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalOneBranchAliasedInputOutputTest) { - // For the conditional branch computation `computation1`, the output is - // aliased with the input, for the other branch, there is no aliasing. The - // aliased input to the conditional is first in the sort order, it should get - // an alternate memory allocation, along with the outputs of both branches, - // and the aliased value custom_call2. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - computation0 { - param0 = (f32[2,2]{1,0}) parameter(0) - gte0 = f32[2,2]{1,0} get-tuple-element(param0), index=0 - negate1 = f32[2,2]{1,0} negate(gte0) - negate2 = f32[2,2]{1,0} negate(negate1) - negate3 = f32[2,2]{1,0} negate(negate2) - negate4 = f32[2,2]{1,0} negate(negate3) - ROOT tuple2 = (f32[2,2]) tuple(negate4) - } - - computation1 { - param1 = (f32[2,2]{1,0}) parameter(0) - gte1 = f32[2,2]{1,0} get-tuple-element(param1), index=0 - negate5 = f32[2,2]{1,0} negate(gte1) - custom_call1 = f32[2,2]{1,0} custom-call(gte1, negate5), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT tuple3 = (f32[2,2]) tuple(custom_call1) - } - - ENTRY entry { - p0 = pred[] parameter(0) - p1 = f32[2,2] parameter(1) - negate0 = f32[2,2]{1,0} negate(p1) - tanh0 = f32[2,2]{1,0} tanh(p1) - tuple0 = (f32[2,2]{1,0}) tuple(negate0) - tuple1 = (f32[2,2]{1,0}) tuple(tanh0) - conditional0 = (f32[2,2]) conditional(p0, tuple0, tuple1), true_computation=computation0, false_computation=computation1 - gte2 = f32[2,2]{1,0} get-tuple-element(conditional0), index=0 - custom_call2 = f32[2,2]{1,0} custom-call(gte2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT negate6 = f32[2,2]{1,0} negate(custom_call2) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 20; - memory_space_options.max_outstanding_prefetches = 0; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "tanh0") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - CheckMemorySpaceForInstructionNames( - module.get(), {"tanh0", "custom_call1", "negate4", "custom_call2"}, - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, ConditionalCommonInputAliasedOutputTest) { - // The input to the conditional is used in both branches, and it is also - // aliased with the custom call output. The custom call output is first in - // sort order, so it gets an alternate memory allocation along with all the - // values that alias with it. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - computation0 { - param0 = (f32[2,2]{1,0}) parameter(0) - gte0 = f32[2,2]{1,0} get-tuple-element(param0), index=0 - negate1 = f32[2,2]{1,0} negate(gte0) - custom_call0 = f32[2,2]{1,0} custom-call(gte0, negate1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT tuple2 = (f32[2,2]) tuple(custom_call0) - } - - computation1 { - param1 = (f32[2,2]{1,0}) parameter(0) - gte1 = f32[2,2]{1,0} get-tuple-element(param1), index=0 - negate2 = f32[2,2]{1,0} negate(gte1) - custom_call1 = f32[2,2]{1,0} custom-call(gte1, negate2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT tuple3 = (f32[2,2]) tuple(custom_call1) - } - - ENTRY entry { - p0 = pred[] parameter(0) - p1 = f32[2,2] parameter(1) - negate0 = f32[2,2]{1,0} negate(p1) - tuple0 = (f32[2,2]{1,0}) tuple(negate0) - conditional0 = (f32[2,2]) conditional(p0, tuple0, tuple0), true_computation=computation0, false_computation=computation1 - gte2 = f32[2,2]{1,0} get-tuple-element(conditional0), index=0 - custom_call2 = f32[2,2]{1,0} custom-call(gte2), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT negate3 = f32[2,2]{1,0} negate(custom_call2) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 1000; - memory_space_options.max_outstanding_prefetches = 0; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "custom_call2") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - CheckMemorySpaceForInstructionNames( - module.get(), {"negate0", "custom_call2", "custom_call0", "custom_call1"}, - kAlternateMemorySpace); -} - -// Tests that a standard async-start/done sequence (without late binding or -// updates) works correctly and places the buffers in alternate memory. -TEST_F(MemorySpaceAssignmentTest, AsyncNoLateBinding) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4] parameter(0) - ROOT custom-call = f32[4] custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4] parameter(0) - negate0 = f32[4] negate(param) - async-start = ((f32[4]), f32[4], s32[]) async-start(negate0), - calls=async_computation - negate1 = f32[4] negate(param) - - async-done = f32[4] async-done(async-start), calls=async_computation - ROOT add = add(async-done, negate1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - HloInstruction* async_start = - module->entry_computation()->GetInstructionWithName("async-start"); - EXPECT_EQ(async_start->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().layout().memory_space(), kAlternateMemorySpace); -} - -// Tests a simple late binding case: async-start has no operands, and the -// operand is bound later via a single async-update. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateBasicLateBinding) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4] parameter(0) - ROOT custom-call = f32[4] custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4] parameter(0) - negate0 = f32[4] negate(param) - async-start = ((), (), s32[]) async-start(), calls=async_computation - negate1 = f32[4] negate(param) - async-update = ((f32[4]), f32[4], s32[]) async-update(async-start, negate0), - calls=async_computation - async-done = f32[4] async-done(async-update), calls=async_computation - ROOT add = add(async-done, negate1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - std::optional options_override = std::move(options); - AssignMemorySpace(module.get(), std::move(options_override)); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - EXPECT_EQ(async_update->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().layout().memory_space(), kAlternateMemorySpace); -} - -// Tests late binding where the operand bound in async-update is defined -// after the async-start instruction. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateLateBindingLateDefinition) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4]{0} parameter(0) - ROOT custom-call = f32[4]{0} custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - async-start = ((), (), s32[]) async-start(), calls=async_computation - negate1 = f32[4]{0} negate(param) - async-update = ((f32[4]{0}), f32[4]{0}, s32[]) async-update(async-start, negate1), - calls=async_computation - async-done = f32[4]{0} async-done(async-update), calls=async_computation - ROOT add = add(async-done, negate1) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - EXPECT_EQ(async_update->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().layout().memory_space(), kAlternateMemorySpace); -} - -// Tests late binding where a parameter is shared between a late-bound operand -// (via an intermediate operation) and another instruction outside the async -// chain. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateLateBindingMultiUseParameter) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4]{0} parameter(0) - ROOT custom-call = f32[4]{0} custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - async-start = ((), (), s32[]) async-start(), calls=async_computation - negate1 = f32[4]{0} negate(param) - async-update = ((f32[4]{0}), f32[4]{0}, s32[]) async-update(async-start, negate1), - calls=async_computation - async-done = f32[4]{0} async-done(async-update), calls=async_computation - ROOT add = add(async-done, param) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - EXPECT_EQ(async_update->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().layout().memory_space(), kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - AsyncUpdateLateBindingOperandDisallowedInAltMem) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4]{0} parameter(0) - ROOT custom-call = f32[4]{0} custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - param2 = f32[4]{0} parameter(1) - async-start = ((), (), s32[]) async-start(), calls=async_computation - negate1 = f32[4]{0} negate(param) - async-update = ((f32[4]{0}), f32[4]{0}, s32[]) async-update(async-start, negate1), - calls=async_computation - async-done = f32[4]{0} async-done(async-update), calls=async_computation - ROOT add = add(async-done, param2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.is_use_allowed_in_alternate_mem_fn = [](const HloUse& use) { - if (use.instruction->opcode() == HloOpcode::kAsyncUpdate && - use.operand_number == 1) { - return false; - } - return true; - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); - std::optional options_override = std::move(options); - AssignMemorySpace(module.get(), std::move(options_override), {}, - &prefetch_interval_picker); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - // The late-bound operand negate1 is restricted to default memory. - EXPECT_EQ(async_update->shape() - .tuple_shapes(0) - .tuple_shapes(0) - .layout() - .memory_space(), - kDefaultMemorySpace); -} - -// Tests late binding with default options, verifying that the late-bound -// operand remains in default memory while the output is placed in alternate -// memory. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateLateBindingMultipleParameters) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4]{0} parameter(0) - ROOT custom-call = f32[4]{0} custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - param2 = f32[4]{0} parameter(1) - async-start = ((), (), s32[]) async-start(), calls=async_computation - negate1 = f32[4]{0} negate(param) - async-update = ((f32[4]{0}), f32[4]{0}, s32[]) async-update(async-start, negate1), - calls=async_computation - async-done = f32[4]{0} async-done(async-update), calls=async_computation - ROOT add = add(async-done, param2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - // Verify the late-bound operand (tuple_shapes(0)) is in default memory space. - EXPECT_EQ(async_update->shape() - .tuple_shapes(0) - .tuple_shapes(0) - .layout() - .memory_space(), - kDefaultMemorySpace); - // Verify the output (tuple_shapes(1)) is in alternate memory space. - EXPECT_EQ(async_update->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().layout().memory_space(), kAlternateMemorySpace); -} - -// Tests late binding where the root instruction uses a parameter, and there -// is another intermediate use of the async output before the root. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateLateBindingChainedUses) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p = f32[4]{0} parameter(0) - ROOT custom-call = f32[4]{0} custom-call(p), custom_call_target="foo" -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - param2 = f32[4]{0} parameter(1) - async-start = ((), (), s32[]) async-start(), calls=async_computation - negate1 = f32[4]{0} negate(param) - async-update = ((f32[4]{0}), f32[4]{0}, s32[]) async-update(async-start, negate1), - calls=async_computation - async-done = f32[4]{0} async-done(async-update), calls=async_computation - add0 = f32[4]{0} add(async-done, param2) - ROOT add = add(add0, param) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - AssignMemorySpace(module.get()); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - EXPECT_EQ(async_update->shape().tuple_shapes(1).layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().layout().memory_space(), kAlternateMemorySpace); -} - -// Tests multiple async-updates in a chain, each adding multiple operands. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateMultipleUpdates) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p0 = f32[4]{0} parameter(0) - p1 = f32[4]{0} parameter(1) - p2 = f32[4]{0} parameter(2) - p3 = f32[4]{0} parameter(3) - ROOT tuple = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(p0, p1, p2, p3) -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - neg0 = f32[4]{0} negate(param) - neg1 = f32[4]{0} negate(param) - neg2 = f32[4]{0} negate(param) - neg3 = f32[4]{0} negate(param) - async-start = ((), (), s32[]) async-start(), calls=async_computation - async-update-1 = ((f32[4]{0}, f32[4]{0}), (), s32[]) async-update(async-start, neg0, neg1), calls=async_computation - async-update-2 = ((f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), s32[]) async-update(async-update-1, neg2, neg3), calls=async_computation - async-done = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) async-done(async-update-2), calls=async_computation - gte0 = f32[4]{0} get-tuple-element(async-done), index=0 - ROOT root = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(gte0, neg0, neg1, neg2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 100); - ASSERT_OK(AssignMemorySpaceAndReturnStatus( - module.get(), /*options_override=*/std::nullopt, - /*buffer_interval_compare=*/std::nullopt, - &prefetch_interval_picker) - .status()); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update-2"); - // Verify the output buffer corresponding to neg0 (first element of output - // tuple of async-update-2) is in default memory space. - // TODO(b/538345137): Update this to check for alternate memory. - EXPECT_EQ(async_update->shape() - .tuple_shapes(1) - .tuple_shapes(0) - .layout() - .memory_space(), - kDefaultMemorySpace); - - // TODO(b/538345137): Update this to check for alternate memory. - HloInstruction* async_done = - module->entry_computation()->GetInstructionWithName("async-done"); - EXPECT_EQ(async_done->shape().tuple_shapes(0).layout().memory_space(), - kDefaultMemorySpace); -} - -// Tests that FindAliases correctly maps operands to parameters for async-start, -// async-update, call, and conditional instructions: -// 1. Operands bound with async-start map 1-to-1 to initial parameters. -// 2. Operand 0 of async-update (the async context token) does not alias -// parameters. -// 3. Operands bound with async-update map 1-to-1 to subsequent parameters, -// offset by the number of previously bound operands. -// 4. Operands of synchronous call map 1-to-1 to callee parameters by operand -// index, and operands 1..N of conditional map to branch_computation(i - 1) -// parameter 0. -TEST_F(MemorySpaceAssignmentTest, FindAliasesAsyncStartAndAsyncUpdate) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p0 = f32[4]{0} parameter(0) - p1 = f32[4]{0} parameter(1) - p2 = f32[4]{0} parameter(2) - p3 = f32[4]{0} parameter(3) - ROOT tuple = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(p0, p1, p2, p3) -} - -call_computation { - cp0 = f32[4]{0} parameter(0) - cp1 = f32[4]{0} parameter(1) - ROOT call_tuple = (f32[4]{0}, f32[4]{0}) tuple(cp0, cp1) -} - -branch0_computation { - ROOT bp0 = f32[4]{0} parameter(0) -} - -branch1_computation { - ROOT bp1 = f32[4]{0} parameter(0) -} - -ENTRY entry { - param0 = f32[4]{0} parameter(0) - param1 = f32[4]{0} parameter(1) - param2 = f32[4]{0} parameter(2) - param3 = f32[4]{0} parameter(3) - param_pred = pred[] parameter(4) - async-start = ((f32[4]{0}, f32[4]{0}), (), s32[]) async-start(param0, param1), calls=async_computation - async-update = ((f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), (), s32[]) async-update(async-start, param2, param3), calls=async_computation - async-done = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) async-done(async-update), calls=async_computation - sync-call = (f32[4]{0}, f32[4]{0}) call(param0, param1), to_apply=call_computation - cond = f32[4]{0} conditional(param_pred, param2, param3), true_computation=branch0_computation, false_computation=branch1_computation - ROOT root = ((f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), (f32[4]{0}, f32[4]{0}), f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(async-done, sync-call, cond, param0, param1, param2, param3) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - HloInstruction* param0 = FindInstruction(module.get(), "param0"); - HloInstruction* param1 = FindInstruction(module.get(), "param1"); - HloInstruction* param2 = FindInstruction(module.get(), "param2"); - HloInstruction* param3 = FindInstruction(module.get(), "param3"); - HloInstruction* pred = FindInstruction(module.get(), "param_pred"); - HloInstruction* async_start = FindInstruction(module.get(), "async-start"); - HloInstruction* async_update = FindInstruction(module.get(), "async-update"); - HloInstruction* async_done = FindInstruction(module.get(), "async-done"); - HloInstruction* sync_call = FindInstruction(module.get(), "sync-call"); - HloInstruction* cond = FindInstruction(module.get(), "cond"); - - HloComputation* async_comp = - module->GetComputationWithName("async_computation"); - HloInstruction* p0 = async_comp->parameter_instruction(0); - HloInstruction* p1 = async_comp->parameter_instruction(1); - HloInstruction* p2 = async_comp->parameter_instruction(2); - HloInstruction* p3 = async_comp->parameter_instruction(3); - - HloComputation* call_comp = - module->GetComputationWithName("call_computation"); - HloInstruction* cp0 = call_comp->parameter_instruction(0); - HloInstruction* cp1 = call_comp->parameter_instruction(1); - - HloComputation* branch0_comp = - module->GetComputationWithName("branch0_computation"); - HloInstruction* bp0 = branch0_comp->parameter_instruction(0); - HloComputation* branch1_comp = - module->GetComputationWithName("branch1_computation"); - HloInstruction* bp1 = branch1_comp->parameter_instruction(0); - - std::vector allocation_values; - // Allocation values for parameters of the async, call, and branch - // computations. - allocation_values.emplace_back(nullptr, HloPosition{p0, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{p1, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{p2, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{p3, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{cp0, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{cp1, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{bp0, {}}, 16); - allocation_values.emplace_back(nullptr, HloPosition{bp1, {}}, 16); - - // Allocation values for caller operands passed to async-start (operands 0 and - // 1) and sync-call (operands 0 and 1). - AllocationValue val_param0(nullptr, HloPosition{param0, {}}, 16); - val_param0.uses().push_back(AllocationValue::Use{HloUse{async_start, 0}, 1}); - val_param0.uses().push_back(AllocationValue::Use{HloUse{sync_call, 0}, 4}); - allocation_values.push_back(std::move(val_param0)); - - AllocationValue val_param1(nullptr, HloPosition{param1, {}}, 16); - val_param1.uses().push_back(AllocationValue::Use{HloUse{async_start, 1}, 1}); - val_param1.uses().push_back(AllocationValue::Use{HloUse{sync_call, 1}, 4}); - allocation_values.push_back(std::move(val_param1)); - - // Allocation values for async-start bound operand elements ({0, 0} and - // {0, 1}) consumed as operand 0 of async-update. - AllocationValue val_async_start_0(nullptr, HloPosition{async_start, {0, 0}}, - 16); - val_async_start_0.uses().push_back( - AllocationValue::Use{HloUse{async_update, 0, {0, 0}}, 2}); - allocation_values.push_back(std::move(val_async_start_0)); - - AllocationValue val_async_start_1(nullptr, HloPosition{async_start, {0, 1}}, - 16); - val_async_start_1.uses().push_back( - AllocationValue::Use{HloUse{async_update, 0, {0, 1}}, 2}); - allocation_values.push_back(std::move(val_async_start_1)); - - // Allocation values for caller operands passed to async-update (operands 1 - // and 2) and conditional (operands 1 and 2). - AllocationValue val_param2(nullptr, HloPosition{param2, {}}, 16); - val_param2.uses().push_back(AllocationValue::Use{HloUse{async_update, 1}, 2}); - val_param2.uses().push_back(AllocationValue::Use{HloUse{cond, 1}, 5}); - allocation_values.push_back(std::move(val_param2)); - - AllocationValue val_param3(nullptr, HloPosition{param3, {}}, 16); - val_param3.uses().push_back(AllocationValue::Use{HloUse{async_update, 2}, 2}); - val_param3.uses().push_back(AllocationValue::Use{HloUse{cond, 2}, 5}); - allocation_values.push_back(std::move(val_param3)); - - AllocationValue val_pred(nullptr, HloPosition{pred, {}}, 1); - val_pred.uses().push_back(AllocationValue::Use{HloUse{cond, 0}, 5}); - allocation_values.push_back(std::move(val_pred)); - - // Allocation values for async-update bound operand elements ({0, 0}, {0, 1}, - // {0, 2}, and {0, 3}) consumed by async-done. - for (int64_t i = 0; i < 4; ++i) { - AllocationValue val_async_update(nullptr, HloPosition{async_update, {0, i}}, - 16); - val_async_update.uses().push_back( - AllocationValue::Use{HloUse{async_done, 0, {0, i}}, 3}); - allocation_values.push_back(std::move(val_async_update)); - } - - MsaAlgorithm::FindAliasesForTesting(&allocation_values); - - auto find_val = [&](const HloInstruction* inst, - ShapeIndex index = {}) -> const AllocationValue& { - for (const auto& v : allocation_values) { - if (v.position().instruction == inst && v.position().index == index) { - return v; - } - } - LOG(FATAL) << "Not found: " << inst->name() << " " << index.ToString(); - }; - - // 1. Verify 1-to-1 mapping for async-start operands: - // param0 (operand 0) aliases {async_start, {0, 0}} and callee parameter p0. - EXPECT_THAT(find_val(param0).uses()[0].aliases, - ::testing::ElementsAre(HloPosition{async_start, {0, 0}}, - HloPosition{p0, {}})); - // param1 (operand 1) aliases {async_start, {0, 1}} and callee parameter p1. - EXPECT_THAT(find_val(param1).uses()[0].aliases, - ::testing::ElementsAre(HloPosition{async_start, {0, 1}}, - HloPosition{p1, {}})); - - // 2. Verify operand 0 of async-update (the async context tuple) aliases - // matching elements in async_update, and does not alias any callee parameter: - EXPECT_THAT(find_val(async_start, {0, 0}).uses()[0].aliases, - ::testing::ElementsAre(HloPosition{async_update, {0, 0}})); - EXPECT_THAT(find_val(async_start, {0, 1}).uses()[0].aliases, - ::testing::ElementsAre(HloPosition{async_update, {0, 1}})); - - // 3. Verify late-bound parameter offset calculation for async-update - // operands: param2 (operand 1 of async-update) aliases {async_update, {0, 2}} - // and callee parameter p2 (1 - 1 + 2 = 2). - EXPECT_THAT(find_val(param2).uses()[0].aliases, - ::testing::ElementsAre(HloPosition{async_update, {0, 2}}, - HloPosition{p2, {}})); - // param3 (operand 2 of async-update) aliases {async_update, {0, 3}} and - // callee parameter p3 (2 - 1 + 2 = 3). - EXPECT_THAT(find_val(param3).uses()[0].aliases, - ::testing::ElementsAre(HloPosition{async_update, {0, 3}}, - HloPosition{p3, {}})); - - // 4. Verify 1-to-1 mapping for synchronous call and conditional operands: - EXPECT_THAT(find_val(param0).uses()[1].aliases, - ::testing::ElementsAre(HloPosition{cp0, {}})); - EXPECT_THAT(find_val(param1).uses()[1].aliases, - ::testing::ElementsAre(HloPosition{cp1, {}})); - EXPECT_THAT(find_val(pred).uses()[0].aliases, ::testing::IsEmpty()); - EXPECT_THAT(find_val(param2).uses()[1].aliases, - ::testing::ElementsAre(HloPosition{bp0, {}})); - EXPECT_THAT(find_val(param3).uses()[1].aliases, - ::testing::ElementsAre(HloPosition{bp1, {}})); -} - -// Tests a case where some operands bound in async-start/update are allowed -// in alternate memory while others are restricted to default memory. -// TODO(b/538345137): Re-enable this test once b/538345137 is fixed. -TEST_F(MemorySpaceAssignmentTest, DISABLED_AsyncUpdateMixedMemorySpaces) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p0 = f32[4]{0} parameter(0) - p1 = f32[4]{0} parameter(1) - p2 = f32[4]{0} parameter(2) - p3 = f32[4]{0} parameter(3) - ROOT tuple = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(p0, p1, p2, p3) -} - -ENTRY entry { - param = f32[4]{0} parameter(0) - neg0 = f32[4]{0} negate(param) - neg1 = f32[4]{0} negate(param) - neg2 = f32[4]{0} negate(param) - neg3 = f32[4]{0} negate(param) - async-start = ((f32[4]{0}, f32[4]{0}), (), s32[]) async-start(neg0, neg1), calls=async_computation - async-update = ((f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}), s32[]) async-update(async-start, neg2, neg3), calls=async_computation - async-done = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) async-done(async-update), calls=async_computation - gte0 = f32[4]{0} get-tuple-element(async-done), index=0 - ROOT root = (f32[4]{0}, f32[4]{0}, f32[4]{0}, f32[4]{0}) tuple(gte0, neg0, neg1, neg2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.is_use_allowed_in_alternate_mem_fn = [](const HloUse& use) { - if (use.instruction->opcode() == HloOpcode::kAsyncStart && - use.operand_number == 1) { - return false; - } - if (use.instruction->opcode() == HloOpcode::kAsyncUpdate && - use.operand_number == 2) { - return false; - } - return true; - }; - std::optional options_override = std::move(options); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 100); - ASSERT_OK(AssignMemorySpaceAndReturnStatus( - module.get(), std::move(options_override), - /*buffer_interval_compare=*/std::nullopt, - &prefetch_interval_picker) - .status()); - - HloInstruction* async_start = - module->entry_computation()->GetInstructionWithName("async-start"); - // neg0 is allowed in alt memory. - EXPECT_EQ(async_start->shape() - .tuple_shapes(0) - .tuple_shapes(0) - .layout() - .memory_space(), - kAlternateMemorySpace); - // neg1 is not allowed in alt memory. - EXPECT_EQ(async_start->shape() - .tuple_shapes(0) - .tuple_shapes(1) - .layout() - .memory_space(), - kDefaultMemorySpace); - - HloInstruction* async_update = - module->entry_computation()->GetInstructionWithName("async-update"); - // neg2 is allowed in alt memory. - EXPECT_EQ(async_update->shape() - .tuple_shapes(0) - .tuple_shapes(2) - .layout() - .memory_space(), - kAlternateMemorySpace); - // neg3 is not allowed in alt memory. - EXPECT_EQ(async_update->shape() - .tuple_shapes(0) - .tuple_shapes(3) - .layout() - .memory_space(), - kDefaultMemorySpace); -} - -// Tests that late-bound operands are correctly handled when repacking is -// triggered due to memory pressure in alternate memory. -TEST_F(MemorySpaceAssignmentTest, AsyncUpdateRepacking) { - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -async_computation { - p0 = f32[2]{0} parameter(0) - ROOT add = f32[2]{0} add(p0, p0) -} - -ENTRY entry { - param = f32[2]{0} parameter(0) - param_large = f32[4]{0} parameter(1) - A = f32[2]{0} negate(param) - D = f32[4]{0} negate(param_large) - use_A = f32[2]{0} negate(A) - use_D = f32[4]{0} negate(D) - async-start = ((), (), s32[]) async-start(), calls=async_computation - add_A = f32[2]{0} add(use_A, param) - add_D = f32[4]{0} add(use_D, param_large) - slice_D = f32[2]{0} slice(add_D), slice={[0:2]} - sum_AD = f32[2]{0} add(add_A, slice_D) - B = f32[2]{0} negate(param) - E = f32[2]{0} negate(param) - use_E = f32[2]{0} negate(E) - async-update = ((f32[2]{0}), f32[2]{0}, s32[]) async-update(async-start, B), calls=async_computation - async-done = f32[2]{0} async-done(async-update), calls=async_computation - add_E = f32[2]{0} add(use_E, param) - sum_1 = f32[2]{0} add(async-done, sum_AD) - ROOT root = f32[2]{0} add(sum_1, add_E) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - // Set memory limit to 24 bytes. - options.max_size_in_bytes = 24; - options.is_use_allowed_in_alternate_mem_fn = [](const HloUse& use) { - if (use.instruction->operand(use.operand_number)->name() == "param" || - use.instruction->operand(use.operand_number)->name() == "param_large") { - return false; - } - absl::string_view name = use.instruction->name(); - if (name == "add_A" || name == "add_D" || name == "slice_D" || - name == "sum_AD" || name == "sum_1" || name == "root") { - return false; - } - return true; - }; - std::optional options_override = std::move(options); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 100); - ASSERT_OK(AssignMemorySpaceAndReturnStatus( - module.get(), std::move(options_override), - /*buffer_interval_compare=*/std::nullopt, - &prefetch_interval_picker) - .status()); - - // Verify all are in alternate memory space 1. - HloInstruction* A_inst = - module->entry_computation()->GetInstructionWithName("A"); - ASSERT_NE(A_inst, nullptr) << "Instruction A not found! Module:\n" - << module->ToString(); - EXPECT_EQ(A_inst->shape().layout().memory_space(), kAlternateMemorySpace); - - HloInstruction* B_inst = - module->entry_computation()->GetInstructionWithName("B"); - ASSERT_NE(B_inst, nullptr) << "Instruction B not found! Module:\n" - << module->ToString(); - EXPECT_EQ(B_inst->shape().layout().memory_space(), kDefaultMemorySpace); - HloInstruction* D_inst = - module->entry_computation()->GetInstructionWithName("D"); - ASSERT_NE(D_inst, nullptr) << "Instruction D not found! Module:\n" - << module->ToString(); - EXPECT_EQ(D_inst->shape().layout().memory_space(), kAlternateMemorySpace); -} - -// Tests that when an allocation value in a computation inherits a non-zero -// computation-level preferred offset, but a specific use (e.g., an aliased -// custom-call output) has a RequiredMemoryAssignment at offset 0, MSA -// successfully overrides the preferred offset with the required assignment's -// offset instead of failing a CHECK_EQ assertion. -TEST_F(MemorySpaceAssignmentTest, - OverridePreferredOffsetWithRequiredAssignmentOffset) { - absl::string_view hlo_string = R"hlo( - HloModule bug, is_scheduled=true - - computation0 { - param0 = (f32[2,2]{1,0}) parameter(0) - gte0 = f32[2,2]{1,0} get-tuple-element(param0), index=0 - negate0 = f32[2,2]{1,0} negate(gte0) - custom_call0 = f32[2,2]{1,0} custom-call(gte0, negate0), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT tuple0 = (f32[2,2]) tuple(custom_call0) - } - - computation1 { - param1 = (f32[2,2]{1,0}) parameter(0) - gte1 = f32[2,2]{1,0} get-tuple-element(param1), index=0 - negate1 = f32[2,2]{1,0} negate(gte1) - custom_call1 = f32[2,2]{1,0} custom-call(gte1, negate1), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT tuple1 = (f32[2,2]) tuple(custom_call1) - } - - ENTRY entry { - p0 = pred[] parameter(0) - p1 = f32[2,2] parameter(1) - negate_entry = f32[2,2]{1,0} negate(p1) - tuple_entry = (f32[2,2]{1,0}) tuple(negate_entry) - conditional = (f32[2,2]) conditional(p0, tuple_entry, tuple_entry), true_computation=computation0, false_computation=computation1 - gte_entry = f32[2,2]{1,0} get-tuple-element(conditional), index=0 - custom_call_entry = f32[2,2]{1,0} custom-call(gte_entry), custom_call_target="tpu_custom_call", output_to_operand_aliasing={{}: (0, {})} - ROOT negate_out = f32[2,2]{1,0} negate(custom_call_entry) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 1000; - memory_space_options.max_outstanding_prefetches = 0; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "custom_call_entry") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - - CheckMemorySpaceForInstructionNames( - module.get(), - {"negate_entry", "custom_call_entry", "custom_call0", "custom_call1"}, - kAlternateMemorySpace); -} - -TEST_F(MemorySpaceAssignmentTest, - ConditionalTestCommonInputColoredUsePostLoop) { - // negate0 is first in sort order, it is used inside the conditional and after - // the conditional. It should get an alternate memory allocation without any - // evictions. The alternate memory is just enough for one of the conditional - // inputs to be in alternate memory. - absl::string_view hlo_string = R"( - HloModule CondAllocation, is_scheduled=true - - computation0 { - param0 = (f32[3,2]{1,0}, f32[3,2]{1,0}) parameter(0) - gte0 = f32[3,2]{1,0} get-tuple-element(param0), index=0 - branch0_gte1 = f32[3,2]{1,0} get-tuple-element(param0), index=1 - branch0_negate = f32[3,2]{1,0} negate(branch0_gte1) - branch0_negate2 = f32[3,2]{1,0} negate(branch0_negate) - branch0_negate3 = f32[3,2]{1,0} negate(branch0_negate2) - branch0_negate4 = f32[3,2]{1,0} negate(branch0_negate3) - branch0_add0 = f32[3,2]{1,0} add(gte0, branch0_negate4) - branch0_negate5 = f32[3,2]{1,0} negate(branch0_add0) - empty0 = f32[2,2]{1,0} custom-call(), custom_call_target="AllocateBuffer" - ROOT tuple2 = (f32[3,2], f32[2,2]) tuple(branch0_negate5, empty0) - } - - computation1 { - param1 = (f32[3,2]{1,0}, f32[3,2]{1,0}) parameter(0) - get1 = f32[3,2]{1,0} get-tuple-element(param1), index=0 - branch1_gte1 = f32[3,2]{1,0} get-tuple-element(param1), index=1 - branch1_negate = f32[3,2]{1,0} negate(branch1_gte1) - branch1_negate2 = f32[3,2]{1,0} negate(branch1_negate) - branch1_negate3 = f32[3,2]{1,0} negate(branch1_negate2) - branch1_negate4 = f32[3,2]{1,0} negate(branch1_negate3) - branch1_add0 = f32[3,2]{1,0} add(get1, branch1_negate4) - branch1_negate5 = f32[3,2]{1,0} negate(branch1_add0) - empty1 = f32[2,2]{1,0} custom-call(), custom_call_target="AllocateBuffer" - ROOT tuple3 = (f32[3,2], f32[2,2]) tuple(branch1_negate5, empty1) - } - - ENTRY entry { - p0 = f32[3,2] parameter(0) - p1 = f32[2,2] parameter(1) - p2 = pred[] parameter(2) - negate0 = f32[3,2]{1,0} negate(p0) - tanh0 = f32[3,2]{1,0} tanh(p0) - tanh1 = f32[3,2]{1,0} tanh(tanh0) - tanh2 = f32[3,2]{1,0} tanh(tanh1) - tuple0 = (f32[3,2]{1,0}, f32[3,2]{1,0}) tuple(negate0, tanh1) - conditional = (f32[3,2], f32[2,2]) conditional(p2, tuple0, tuple0), true_computation=computation0, false_computation=computation1 - gte2 = f32[3,2]{1,0} get-tuple-element(conditional), index=0 - gte3 = f32[2,2]{1,0} get-tuple-element(conditional), index=1 - add6 = f32[3,2]{1,0} add(negate0, gte2) - add7 = f32[2,2]{1,0} add(gte3, p1) - ROOT tuple4 = (f32[3,2]{1,0}, f32[2,2]{1,0}) tuple(add6, add7) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.max_outstanding_prefetches = 2; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - // An arbitrary value that is greater than that for p1, p2, p3, and - // p4. - int priority = 100; - if (x.buffer->instruction()->name() == "negate0") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - CheckMemorySpaceForInstructionNames(module.get(), {"negate0"}, - kAlternateMemorySpace); - - // Check that the operand 0 of instructions add6 and tuple0 are in alternate - // memory. Check that they are not copy done instructions. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"tuple0", "add6"}, 0, kAlternateMemorySpace, - HloOpcode::kNegate); -} - -TEST_F(MemorySpaceAssignmentTest, - NestedConditionalInputPreAllocatedInAltMemory) { - // The input, tanh0, to the outer conditional is pre-allocated in alternate - // memory. The same input is also used in the nested conditional. All - // positions of tanh0 in both conditionals should in alternate memory. - absl::string_view hlo_string = R"hlo( - HloModule CondAllocation, is_scheduled=true - - true_computation2 { - p0_tc2 = (f32[3]{0}) parameter(0) - gte_tc2 = f32[3]{0} get-tuple-element(p0_tc2), index=0 - ROOT neg1 = f32[3]{0} negate(gte_tc2) - } - - false_computation2 { - p0_fc2 = (f32[3]{0}) parameter(0) - gte_fc2 = f32[3]{0} get-tuple-element(p0_fc2), index=0 - copy = f32[3]{0} copy(gte_fc2) - ROOT neg2 = f32[3]{0} negate(copy) - } - - true_computation1 { - p0_tc1 = (f32[3]{0}) parameter(0) - gte_tc1 = f32[3]{0} get-tuple-element(p0_tc1), index=0 - slice = f32[1]{0} slice(gte_tc1), slice={[0:1]} - bitcast = f32[] bitcast(slice) - constant = f32[] constant(0.0) - compare = pred[] compare(bitcast, constant), direction=GT - conditional_tc1 = f32[3]{0} conditional(compare, p0_tc1, p0_tc1), true_computation=true_computation2, false_computation=false_computation2 - ROOT tuple_tc1 = (f32[3]{0}) tuple(conditional_tc1) - } - - false_computation1 { - p0_fc1 = (f32[3]{0}) parameter(0) - gte_fc1 = f32[3]{0} get-tuple-element(p0_fc1), index=0 - neg3 = f32[3]{0} negate(gte_fc1) - ROOT tuple_fc1 = (f32[3]{0}) tuple(neg3) - } - - - ENTRY entry { - p0 = f32[3]{0} parameter(0) - p1 = pred[] parameter(1) - tanh0 = f32[3]{0} tanh(p0) - tuple = (f32[3]{0}) tuple(tanh0) - conditional = (f32[3]{0}) conditional(p1, tuple, tuple), true_computation=true_computation1, false_computation=false_computation1 - gte = f32[3]{0} get-tuple-element(conditional), index=0 - ROOT add0 = f32[3]{0} add(tanh0, gte) - } - )hlo"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.max_outstanding_prefetches = 2; - memory_space_options.enable_sync_copy_replacement = true; - memory_space_options.enable_sync_slice_replacement = true; - MsaBufferIntervalCompare buffer_interval_compare = - [](const MsaBufferInterval& lhs, const MsaBufferInterval& rhs) { - auto lookup = [](const MsaBufferInterval& x) { - int priority = 100; - if (x.buffer->instruction()->name() == "tanh0") { - priority = 1; - } - return std::make_tuple(priority, x.buffer->instruction()->name()); - }; - - return lookup(lhs) < lookup(rhs); - }; - InstructionCountPrefetchIntervalPicker prefetch_interval_picker(0, 1000); - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpace(module.get(), std::move(memory_space_options), - buffer_interval_compare, &prefetch_interval_picker); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // Check all positions of tanh0 are in alternate memory. - CheckMemorySpaceForInstructionNames( - module.get(), {"tanh0", "gte_tc1", "gte_fc1", "gte_tc2", "gte_fc2"}, - kAlternateMemorySpace); - - // Check that the operand 0 of instructions add0 and tuple are in alternate - // memory. Verify that they read from tanh0 and not from an evicted copy. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - module.get(), {"tuple", "add0"}, 0, kAlternateMemorySpace, - HloOpcode::kTanh); -} - -TEST_F(MemorySpaceAssignmentTest, TestBlockPrefetchingBitcastUses) { - // Test block prefetching with two program inputs that can be block prefetched - // sequentially if bitcasts are not considered as an HLO use when computing - // first use times. - // - // p0 and p1 have size 24 bytes (equal to block prefetch limit of 24 bytes). - // p2 has size 48 bytes (larger than block prefetch limit of 24 bytes). - // bitcast(p0) and bitcast(p1) are scheduled early at the beginning. - // - // Because bitcasts are no-ops and skipped when finding first use times: - // - p2 is not prefetched because its size exceeds the limit. - // - p1 is prefetched and used at negate_p1_0. - // - p0 is prefetched after p1's live range ends and used at negate_p0_0. - absl::string_view hlo_string = R"hlo( -HloModule module, is_scheduled=true - -ENTRY entry { - p0 = f32[2,3]{1,0} parameter(0) - p1 = f32[2,3]{1,0} parameter(1) - p2 = f32[4,3]{1,0} parameter(2) - bitcast_p0 = f32[6]{0} bitcast(p0) - bitcast_p1 = f32[6]{0} bitcast(p1) - negate_p2_0 = f32[4,3]{1,0} negate(p2) - negate_p2_1 = f32[4,3]{1,0} negate(negate_p2_0) - negate_p2_2 = f32[4,3]{1,0} negate(negate_p2_1) - negate_p1_0 = f32[6]{0} negate(bitcast_p1) - negate_p1_1 = f32[6]{0} negate(negate_p1_0) - negate_p1_2 = f32[6]{0} negate(negate_p1_1) - negate_p0_0 = f32[6]{0} negate(bitcast_p0) - negate_p0_1 = f32[6]{0} negate(negate_p0_0) - negate_p0_2 = f32[6]{0} negate(negate_p0_1) - ROOT tuple = (f32[4,3]{1,0}, f32[6]{0}, f32[6]{0}) tuple(negate_p2_2, negate_p1_2, negate_p0_2) -})hlo"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - Options memory_space_options = DefaultMemorySpaceOptions(); - memory_space_options.max_size_in_bytes = 24; - memory_space_options.reserved_bytes_for_block_prefetches = 24; - memory_space_options.max_outstanding_block_prefetches = 10; - memory_space_options.max_outstanding_prefetches = 0; - - memory_space_options.block_prefetched_positions = GetHloPositions( - /*module=*/module.get(), - /*instruction_names=*/{"p0", "p1", "p2"}); - - XLA_VLOG_LINES(1, "Before MSA: \n" + module->ToString()); - AssignMemorySpaceUsingCostAnalysis(module.get(), - std::move(memory_space_options)); - XLA_VLOG_LINES(1, "After MSA: \n" + module->ToString()); - - // p2 exceeds the block prefetch limit, so negate_p2_0 reads p2 from default - // memory. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), /*instruction_names=*/{"negate_p2_0"}, - /*operand_number=*/0, /*operand_memory_space=*/kDefaultMemorySpace, - /*operand_opcode=*/HloOpcode::kParameter); - - // negate_p0_0 and negate_p1_0 should read from bitcast in alternate memory. - CheckOperandOpcodeAndMemorySpaceForInstructionNames( - /*module=*/module.get(), - /*instruction_names=*/{"negate_p0_0", "negate_p1_0"}, - /*operand_number=*/0, /*operand_memory_space=*/kAlternateMemorySpace, - /*operand_opcode=*/HloOpcode::kBitcast); - - // Verify that the bitcast operands read from CopyDone in alternate memory. - HloInstruction* negate_p1_0 = FindInstruction(module.get(), "negate_p1_0"); - ASSERT_NE(negate_p1_0, nullptr); - const HloInstruction* bitcast_p1 = negate_p1_0->operand(0); - EXPECT_EQ(bitcast_p1->opcode(), HloOpcode::kBitcast); - EXPECT_EQ(bitcast_p1->shape().layout().memory_space(), kAlternateMemorySpace); - const HloInstruction* copy_done_p1 = bitcast_p1->operand(0); - EXPECT_EQ(copy_done_p1->opcode(), HloOpcode::kCopyDone); - EXPECT_EQ(copy_done_p1->shape().layout().memory_space(), - kAlternateMemorySpace); - - HloInstruction* negate_p0_0 = FindInstruction(module.get(), "negate_p0_0"); - ASSERT_NE(negate_p0_0, nullptr); - const HloInstruction* bitcast_p0 = negate_p0_0->operand(0); - EXPECT_EQ(bitcast_p0->opcode(), HloOpcode::kBitcast); - EXPECT_EQ(bitcast_p0->shape().layout().memory_space(), kAlternateMemorySpace); - const HloInstruction* copy_done_p0 = bitcast_p0->operand(0); - EXPECT_EQ(copy_done_p0->opcode(), HloOpcode::kCopyDone); - EXPECT_EQ(copy_done_p0->shape().layout().memory_space(), - kAlternateMemorySpace); -} - -// Returns OK if all the following conditions are met: -// 1. All positions belong to the same HloBuffer. -// 2. All positions are found in preset_assignments. -// 3. All positions have the same offset in preset_assignments. -absl::Status ArePositionsColocatedInAltMem( - absl::Span positions, - const PresetAssignments* preset_assignments, - const HloAliasAnalysis* alias_analysis) { - if (positions.empty()) { - return absl::OkStatus(); - } - - // Check that all positions belong to the same HloBuffer. - std::optional buffer_id; - for (const HloPosition& position : positions) { - const HloBuffer& buffer = - alias_analysis->GetUniqueBufferAt(position.instruction, position.index); - if (!buffer_id.has_value()) { - buffer_id = buffer.id(); - } else if (buffer_id.value() != buffer.id()) { - return absl::InternalError(absl::StrCat( - "Positions do not belong to the same HloBuffer. Expected buffer id ", - buffer_id.value(), " but position ", position.instruction->name(), - " index ", position.index.ToString(), " has buffer id ", buffer.id(), - ".")); - } - } - - // Check that all positions in preset_assignments have the same offset. - std::optional offset; - const HloPosition* first_chunk_pos = nullptr; - int chunk_found_count = 0; - for (const HloPosition& position : positions) { - for (const auto& chunk : preset_assignments->chunks()) { - if (chunk.first.instruction == position.instruction && - chunk.first.index == position.index) { - chunk_found_count++; - if (!offset.has_value()) { - offset = chunk.second.offset; - first_chunk_pos = &position; - } else if (offset.value() != chunk.second.offset) { - return absl::InternalError(absl::StrCat( - "Offsets do not match. Position ", - first_chunk_pos->instruction->name(), " index ", - first_chunk_pos->index.ToString(), " has offset ", offset.value(), - ", but position ", position.instruction->name(), " index ", - position.index.ToString(), " has offset ", chunk.second.offset, - ".")); - } - break; - } - } - } - - // Check that at least one position was found in preset_assignments. - if (chunk_found_count == 0) { - return absl::InternalError( - "None of the provided positions were found in preset_assignments."); - } - - return absl::OkStatus(); -} - -TEST_F(MemorySpaceAssignmentTest, - AsyncPipelinedWhileLoopAsyncSliceColocationAndIntermediateVmemReuse) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -ds_comp1 { - p_base = f32[8,16]{1,0} parameter(0) - p_idx0 = s32[] parameter(1) - p_idx1 = s32[] parameter(2) - ROOT ds = f32[2,16]{1,0} dynamic-slice(p_base, p_idx0, p_idx1), dynamic_slice_sizes={2,16} -} - -ds_comp2 { - p_base = f32[8,16]{1,0} parameter(0) - p_idx0 = s32[] parameter(1) - p_idx1 = s32[] parameter(2) - ROOT ds = f32[2,16]{1,0} dynamic-slice(p_base, p_idx0, p_idx1), dynamic_slice_sizes={2,16} -} - -while_cond { - param = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) parameter(0) - i = s32[] get-tuple-element(param), index=0 - limit = s32[] constant(4) - ROOT cond = pred[] compare(i, limit), direction=LT -} - -while_body { - param = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) parameter(0) - i = s32[] get-tuple-element(param), index=0 - async_in1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(param), index=1 - async_in2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(param), index=2 - acc = f32[2,16]{1,0} get-tuple-element(param), index=3 - base1 = f32[8,16]{1,0} get-tuple-element(param), index=4 - base2 = f32[8,16]{1,0} get-tuple-element(param), index=5 - one = s32[] constant(1) - next_i = s32[] add(i, one) - ds_done1 = f32[2,16]{1,0} async-done(async_in1), calls=ds_comp1 - ds_done2 = f32[2,16]{1,0} async-done(async_in2), calls=ds_comp2 - new_acc1 = f32[2,16]{1,0} add(acc, ds_done1) - new_acc2 = f32[2,16]{1,0} add(new_acc1, ds_done2) - new_acc3 = f32[2,16]{1,0} add(new_acc2, new_acc2) - new_acc4 = f32[2,16]{1,0} add(new_acc3, new_acc3) - zero = s32[] constant(0) - next_row = s32[] multiply(next_i, s32[] constant(2)) - async_next1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(base1, next_row, zero), calls=ds_comp1 - async_next2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(base2, next_row, zero), calls=ds_comp2 - ROOT next_state = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) tuple(next_i, async_next1, async_next2, new_acc4, base1, base2) -} - -ENTRY main { - p_base1 = f32[8,16]{1,0} parameter(0) - p_base2 = f32[8,16]{1,0} parameter(1) - c0 = s32[] constant(0) - c_idx0 = s32[] constant(0) - c_idx1 = s32[] constant(0) - c_zero = f32[] constant(0.0) - init_acc = f32[2,16]{1,0} broadcast(c_zero), dimensions={} - prologue_start1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(p_base1, c_idx0, c_idx1), calls=ds_comp1 - prologue_start2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(p_base2, c_idx0, c_idx1), calls=ds_comp2 - init = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) tuple(c0, prologue_start1, prologue_start2, init_acc, p_base1, p_base2) - loop = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) while(init), condition=while_cond, body=while_body, frontend_attributes={xla_disable_while_loop_copies="true"} - loop_async_out1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(loop), index=1 - loop_async_out2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(loop), index=2 - loop_acc_out = f32[2,16]{1,0} get-tuple-element(loop), index=3 - epilogue_done1 = f32[2,16]{1,0} async-done(loop_async_out1), calls=ds_comp1 - epilogue_done2 = f32[2,16]{1,0} async-done(loop_async_out2), calls=ds_comp2 - acc_with_done1 = f32[2,16]{1,0} add(loop_acc_out, epilogue_done1) - ROOT acc_with_done2 = f32[2,16]{1,0} add(acc_with_done1, epilogue_done2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 256; - options.max_outstanding_prefetches = 0; - - MsaBufferIntervalCompare buffer_interval_compare = - CreateBufferIntervalCompareFnFromInstructionNames({"new_acc3"}); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker( - 2, /*max_overlap_count=*/10000); - - auto preset_assignments = - AssignMemorySpace(module.get(), std::move(options), - buffer_interval_compare, &prefetch_interval_picker); - - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - - HloInstruction* while_instr = - module->entry_computation()->GetInstructionWithName("loop"); - ASSERT_NE(while_instr, nullptr); - HloComputation* while_body = while_instr->while_body(); - HloInstruction* body_param = while_body->parameter_instruction(0); - HloInstruction* prologue_start1 = - module->entry_computation()->GetInstructionWithName("prologue_start1"); - HloInstruction* prologue_start2 = - module->entry_computation()->GetInstructionWithName("prologue_start2"); - HloInstruction* epilogue_done1 = - module->entry_computation()->GetInstructionWithName("epilogue_done1"); - HloInstruction* epilogue_done2 = - module->entry_computation()->GetInstructionWithName("epilogue_done2"); - HloInstruction* async_next1 = - while_body->GetInstructionWithName("async_next1"); - HloInstruction* async_next2 = - while_body->GetInstructionWithName("async_next2"); - HloInstruction* new_acc3 = while_body->GetInstructionWithName("new_acc3"); - - // Verify that all of the positions for the first pipelined async slice are - // colocated in alternate memory. - EXPECT_OK(ArePositionsColocatedInAltMem( - { - HloPosition{prologue_start1, {1}}, - HloPosition{while_instr, {1, 1}}, - HloPosition{epilogue_done1, {}}, - HloPosition{body_param, {1, 1}}, - HloPosition{async_next1, {1}}, - }, - preset_assignments.get(), alias_analysis.get())); - - // Verify that all of the positions for the second pipelined async slice are - // colocated in alternate memory. - EXPECT_OK(ArePositionsColocatedInAltMem( - { - HloPosition{prologue_start2, {1}}, - HloPosition{while_instr, {2, 1}}, - HloPosition{epilogue_done2, {}}, - HloPosition{body_param, {2, 1}}, - HloPosition{async_next2, {1}}, - }, - preset_assignments.get(), alias_analysis.get())); - - // Verify that new_acc3 gets alt mem. This demonstrates that we can reuse the - // alt mem allocated for the pipelined async slices, when they are not alive - // in the middle of the while loop. - EXPECT_OK(ArePositionsColocatedInAltMem( - { - HloPosition{new_acc3, {}}, - }, - preset_assignments.get(), alias_analysis.get())); -} - -TEST_F(MemorySpaceAssignmentTest, - AsyncPipelinedWhileLoopColocationWithPrefetchesEnabled) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -ds_comp1 { - p_base = f32[8,16]{1,0} parameter(0) - p_idx0 = s32[] parameter(1) - p_idx1 = s32[] parameter(2) - ROOT ds = f32[2,16]{1,0} dynamic-slice(p_base, p_idx0, p_idx1), dynamic_slice_sizes={2,16} -} - -ds_comp2 { - p_base = f32[8,16]{1,0} parameter(0) - p_idx0 = s32[] parameter(1) - p_idx1 = s32[] parameter(2) - ROOT ds = f32[2,16]{1,0} dynamic-slice(p_base, p_idx0, p_idx1), dynamic_slice_sizes={2,16} -} - -while_cond { - param = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) parameter(0) - i = s32[] get-tuple-element(param), index=0 - limit = s32[] constant(4) - ROOT cond = pred[] compare(i, limit), direction=LT -} - -while_body { - param = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) parameter(0) - i = s32[] get-tuple-element(param), index=0 - async_in1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(param), index=1 - async_in2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(param), index=2 - acc = f32[2,16]{1,0} get-tuple-element(param), index=3 - base1 = f32[8,16]{1,0} get-tuple-element(param), index=4 - base2 = f32[8,16]{1,0} get-tuple-element(param), index=5 - one = s32[] constant(1) - next_i = s32[] add(i, one) - ds_done1 = f32[2,16]{1,0} async-done(async_in1), calls=ds_comp1 - ds_done2 = f32[2,16]{1,0} async-done(async_in2), calls=ds_comp2 - new_acc1 = f32[2,16]{1,0} add(acc, ds_done1) - new_acc2 = f32[2,16]{1,0} add(new_acc1, ds_done2) - zero = s32[] constant(0) - next_row = s32[] multiply(next_i, s32[] constant(2)) - async_next1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(base1, next_row, zero), calls=ds_comp1 - async_next2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(base2, next_row, zero), calls=ds_comp2 - ROOT next_state = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) tuple(next_i, async_next1, async_next2, new_acc2, base1, base2) -} - -ENTRY main { - p_base1 = f32[8,16]{1,0} parameter(0) - p_base2 = f32[8,16]{1,0} parameter(1) - c0 = s32[] constant(0) - c_idx0 = s32[] constant(0) - c_idx1 = s32[] constant(0) - c_zero = f32[] constant(0.0) - init_acc = f32[2,16]{1,0} broadcast(c_zero), dimensions={} - prologue_start1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(p_base1, c_idx0, c_idx1), calls=ds_comp1 - prologue_start2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(p_base2, c_idx0, c_idx1), calls=ds_comp2 - init = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) tuple(c0, prologue_start1, prologue_start2, init_acc, p_base1, p_base2) - loop = (s32[], ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]), f32[2,16]{1,0}, f32[8,16]{1,0}, f32[8,16]{1,0}) while(init), condition=while_cond, body=while_body, frontend_attributes={xla_disable_while_loop_copies="true"} - loop_async_out1 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(loop), index=1 - loop_async_out2 = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) get-tuple-element(loop), index=2 - loop_acc_out = f32[2,16]{1,0} get-tuple-element(loop), index=3 - epilogue_done1 = f32[2,16]{1,0} async-done(loop_async_out1), calls=ds_comp1 - epilogue_done2 = f32[2,16]{1,0} async-done(loop_async_out2), calls=ds_comp2 - acc_with_done1 = f32[2,16]{1,0} add(loop_acc_out, epilogue_done1) - ROOT acc_with_done2 = f32[2,16]{1,0} add(acc_with_done1, epilogue_done2) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 512; - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker( - 2, /*max_overlap_count=*/10000); - - auto preset_assignments = AssignMemorySpace( - module.get(), std::move(options), - /*buffer_interval_compare=*/std::nullopt, &prefetch_interval_picker); - - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - - HloInstruction* while_instr = - module->entry_computation()->GetInstructionWithName("loop"); - ASSERT_NE(while_instr, nullptr); - HloComputation* while_body = while_instr->while_body(); - HloInstruction* body_param = while_body->parameter_instruction(0); - HloInstruction* prologue_start1 = - module->entry_computation()->GetInstructionWithName("prologue_start1"); - HloInstruction* prologue_start2 = - module->entry_computation()->GetInstructionWithName("prologue_start2"); - HloInstruction* epilogue_done1 = - module->entry_computation()->GetInstructionWithName("epilogue_done1"); - HloInstruction* epilogue_done2 = - module->entry_computation()->GetInstructionWithName("epilogue_done2"); - HloInstruction* async_next1 = - while_body->GetInstructionWithName("async_next1"); - HloInstruction* async_next2 = - while_body->GetInstructionWithName("async_next2"); - - // Verify that all of the positions for the first pipelined async slice are - // colocated in alternate memory. - EXPECT_OK(ArePositionsColocatedInAltMem( - { - HloPosition{prologue_start1, {1}}, - HloPosition{while_instr, {1, 1}}, - HloPosition{epilogue_done1, {}}, - HloPosition{body_param, {1, 1}}, - HloPosition{async_next1, {1}}, - }, - preset_assignments.get(), alias_analysis.get())); - - // Verify that all of the positions for the second pipelined async slice are - // colocated in alternate memory. - EXPECT_OK(ArePositionsColocatedInAltMem( - { - HloPosition{prologue_start2, {1}}, - HloPosition{while_instr, {2, 1}}, - HloPosition{epilogue_done2, {}}, - HloPosition{body_param, {2, 1}}, - HloPosition{async_next2, {1}}, - }, - preset_assignments.get(), alias_analysis.get())); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncStartAcrossUntouchedWhileLoopTest) { - absl::string_view hlo_string = R"hlo( -HloModule Module, is_scheduled=true - -ds_comp { - p_base = f32[8,16]{1,0} parameter(0) - p_idx0 = s32[] parameter(1) - p_idx1 = s32[] parameter(2) - ROOT ds = f32[2,16]{1,0} dynamic-slice(p_base, p_idx0, p_idx1), dynamic_slice_sizes={2,16} -} - -while_cond { - param = s32[] parameter(0) - limit = s32[] constant(4) - ROOT cond = pred[] compare(param, limit), direction=LT -} - -while_body { - param = s32[] parameter(0) - one = s32[] constant(1) - ROOT next_param = s32[] add(param, one) -} - -ENTRY main { - p_base = f32[8,16]{1,0} parameter(0) - c0 = s32[] constant(0) - c_idx0 = s32[] constant(0) - c_idx1 = s32[] constant(0) - start = ((f32[8,16]{1,0}, s32[], s32[]), f32[2,16]{1,0}, s32[]) async-start(p_base, c_idx0, c_idx1), calls=ds_comp - loop = s32[] while(c0), condition=while_cond, body=while_body - done = f32[2,16]{1,0} async-done(start), calls=ds_comp - ROOT result = (s32[], f32[2,16]{1,0}) tuple(loop, done) -} - )hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - Options options = DefaultMemorySpaceOptions(); - options.max_size_in_bytes = 128; - - InstructionCountPrefetchIntervalPicker prefetch_interval_picker( - 2, /*max_overlap_count=*/10000); - - auto preset_assignments = AssignMemorySpace( - module.get(), std::move(options), - /*buffer_interval_compare=*/std::nullopt, &prefetch_interval_picker); - - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - - HloInstruction* start = - module->entry_computation()->GetInstructionWithName("start"); - HloInstruction* done = - module->entry_computation()->GetInstructionWithName("done"); - - EXPECT_OK(ArePositionsColocatedInAltMem( - { - HloPosition{start, {1}}, - HloPosition{done, {}}, - }, - preset_assignments.get(), alias_analysis.get())); -} - -TEST_F(MemorySpaceAssignmentTest, AsyncBarrierAliasingCycle) { - absl::string_view hlo_string = R"hlo( -HloModule sc_async_barrier_aliasing, is_scheduled=true - - sc_custom_call_comp { - ROOT %custom_call = (f32[32,128,8]{1,2,0}, u32[], u32[]) custom-call(), - custom_call_target="SparseCoreBarrierStart", - custom_call_has_side_effect=true - } - - sc_main { - %input = f32[32,128,8]{1,2,0} parameter(0) - %output = f32[32,128,8]{1,2,0} parameter(1) - %send_sflag = u32[] parameter(2) - %recv_sflag = u32[] parameter(3) - ROOT %a2a-result = f32[32,128,8]{1,2,0} custom-call(%input, %output, %send_sflag, %recv_sflag), - custom_call_target="SparseCoreBarrierDoneAndCollective", - output_to_operand_aliasing={{}: (1, {})} - } - -ENTRY %Comp_spmd { - %p0 = f32[32,128,8]{1,2,0} parameter(0) - %p1 = f32[32,128,8]{1,2,0} parameter(1) - - %sc_custom_call_start = ((), (f32[32,128,8]{1,2,0}, u32[], u32[]), s32[]) call-start(), - to_apply=%sc_custom_call_comp - %sc_custom_call_done = (f32[32,128,8]{1,2,0}, u32[], u32[]) call-done(%sc_custom_call_start) - %send_sflag = u32[] get-tuple-element(%sc_custom_call_done), index=1 - %recv_sflag = u32[] get-tuple-element(%sc_custom_call_done), index=2 - - %sc_start = ((f32[32,128,8]{1,2,0}, f32[32,128,8]{1,2,0}, u32[], u32[]), f32[32,128,8]{1,2,0}, s32[]) - call-start(%p0, %p1, %send_sflag, %recv_sflag), - to_apply=%sc_main, - output_to_operand_aliasing={ - {0,1}: (1, {}), // update output buffer at {0,0} - {0,2}: (2, {}), // update send_sflag - {0,3}: (3, {}), // update recv_sflag - {1}: (1, {})} // update output buffer at {1} - %sc_output = f32[32,128,8]{1,2,0} call-done(%sc_start) - %neg = f32[32,128,8]{1,2,0} negate(%sc_output) - - ROOT %root = (f32[32,128,8]{1,2,0}, f32[32,128,8]{1,2,0}) tuple(%neg, %p0) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - EXPECT_NO_FATAL_FAILURE(AssignMemorySpace(module.get())); -} - -TEST_F(MemorySpaceAssignmentTest, - AsyncPipelinedWhileAlternateMemoryPositionSplashAttention) { - // Realistic HLO pattern from Splash Attention forward pass with async - // pipelining enabled (xla_disable_while_loop_copies="true"). - // DynamicSlice prefetches Key/Value slices from large HBM base buffers into - // alternate memory (VMEM). - // DynamicUpdateSlice writes Output accumulator slices from VMEM into a large - // HBM base buffer. - const absl::string_view hlo_string = R"hlo( -HloModule SplashAttentionPipelinedWhile, is_scheduled=true - -%async_ds_comp (base: bf16[4,384,128], i0: s32[], i1: s32[], i2: s32[]) -> bf16[1,128,128] { - %base = bf16[4,384,128]{2,1,0} parameter(0) - %i0 = s32[] parameter(1) - %i1 = s32[] parameter(2) - %i2 = s32[] parameter(3) - ROOT %ds = bf16[1,128,128]{2,1,0} dynamic-slice(%base, %i0, %i1, %i2), dynamic_slice_sizes={1,128,128} -} - -%async_dus_comp (base: bf16[4,384,128], update: bf16[1,128,128], i0: s32[], i1: s32[], i2: s32[]) -> bf16[4,384,128] { - %base = bf16[4,384,128]{2,1,0} parameter(0) - %update = bf16[1,128,128]{2,1,0} parameter(1) - %i0 = s32[] parameter(2) - %i1 = s32[] parameter(3) - %i2 = s32[] parameter(4) - ROOT %dus = bf16[4,384,128]{2,1,0} dynamic-update-slice(%base, %update, %i0, %i1, %i2) -} - -%WhileBody (body_param: (s32[], ((bf16[4,384,128], s32[], s32[], s32[]), bf16[1,128,128], s32[]), ((bf16[4,384,128], bf16[1,128,128], s32[], s32[], s32[]), bf16[4,384,128], s32[]))) -> (s32[], ((bf16[4,384,128], s32[], s32[], s32[]), bf16[1,128,128], s32[]), ((bf16[4,384,128], bf16[1,128,128], s32[], s32[], s32[]), bf16[4,384,128], s32[])) { - %body_param = (s32[], ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]), ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[])) parameter(0) - %iter = s32[] get-tuple-element(%body_param), index=0 - %c1 = s32[] constant(1) - %next_iter = s32[] add(%iter, %c1) - %ds_state = ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]) get-tuple-element(%body_param), index=1 - %slice_data = bf16[1,128,128]{2,1,0} async-done(%ds_state), calls=%async_ds_comp - %dus_state = ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[]) get-tuple-element(%body_param), index=2 - %dus_done = bf16[4,384,128]{2,1,0} async-done(%dus_state), calls=%async_dus_comp - %ds_contexts = (bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]) get-tuple-element(%ds_state), index=0 - %k_base = bf16[4,384,128]{2,1,0} get-tuple-element(%ds_contexts), index=0 - %c0 = s32[] constant(0) - %next_ds = ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]) async-start(%k_base, %c0, %c0, %c0), calls=%async_ds_comp - %next_dus = ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[]) async-start(%dus_done, %slice_data, %c0, %c0, %c0), output_to_operand_aliasing={{1}: (0, {})}, calls=%async_dus_comp - ROOT %root_tuple = (s32[], ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]), ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[])) tuple(%next_iter, %next_ds, %next_dus) -} - -%WhileCond (cond_param: (s32[], ((bf16[4,384,128], s32[], s32[], s32[]), bf16[1,128,128], s32[]), ((bf16[4,384,128], bf16[1,128,128], s32[], s32[], s32[]), bf16[4,384,128], s32[]))) -> pred[] { - %cond_param = (s32[], ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]), ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[])) parameter(0) - %iter = s32[] get-tuple-element(%cond_param), index=0 - %limit = s32[] constant(16) - ROOT %cmp = pred[] compare(%iter, %limit), direction=LT -} - -ENTRY %Entry (k_base: bf16[4,384,128], o_base: bf16[4,384,128], update_slice: bf16[1,128,128]) -> (bf16[1,128,128], bf16[4,384,128]) { - %k_base = bf16[4,384,128]{2,1,0} parameter(0) - %o_base = bf16[4,384,128]{2,1,0} parameter(1) - %update_slice = bf16[1,128,128]{2,1,0} parameter(2) - %c0 = s32[] constant(0) - %init_ds = ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]) async-start(%k_base, %c0, %c0, %c0), calls=%async_ds_comp - %init_dus = ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[]) async-start(%o_base, %update_slice, %c0, %c0, %c0), output_to_operand_aliasing={{1}: (0, {})}, calls=%async_dus_comp - %init_tuple = (s32[], ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]), ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[])) tuple(%c0, %init_ds, %init_dus) - %while = (s32[], ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]), ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[])) while(%init_tuple), condition=%WhileCond, body=%WhileBody, frontend_attributes={xla_disable_while_loop_copies="true"} - %final_ds_state = ((bf16[4,384,128]{2,1,0}, s32[], s32[], s32[]), bf16[1,128,128]{2,1,0}, s32[]) get-tuple-element(%while), index=1 - %final_slice = bf16[1,128,128]{2,1,0} async-done(%final_ds_state), calls=%async_ds_comp - %final_dus_state = ((bf16[4,384,128]{2,1,0}, bf16[1,128,128]{2,1,0}, s32[], s32[], s32[]), bf16[4,384,128]{2,1,0}, s32[]) get-tuple-element(%while), index=2 - %final_dus = bf16[4,384,128]{2,1,0} async-done(%final_dus_state), calls=%async_dus_comp - ROOT %out_tuple = (bf16[1,128,128]{2,1,0}, bf16[4,384,128]{2,1,0}) tuple(%final_slice, %final_dus) -} -)hlo"; - - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - - HloInstruction* while_inst = - module->entry_computation()->GetInstructionWithName("while"); - ASSERT_NE(while_inst, nullptr); - HloInstruction* body_param = - while_inst->while_body()->parameter_instruction(0); - ASSERT_NE(body_param, nullptr); - HloInstruction* root_tuple = while_inst->while_body()->root_instruction(); - ASSERT_NE(root_tuple, nullptr); - - // 1. Verify wrapped opcode inspection for the pipelined while loop: - EXPECT_EQ(GetAsyncPipelinedWhileWrappedOpcode(while_inst, /*tuple_idx=*/1, - *alias_analysis), - HloOpcode::kDynamicSlice); - EXPECT_EQ(GetAsyncPipelinedWhileWrappedOpcode(while_inst, /*tuple_idx=*/2, - *alias_analysis), - HloOpcode::kDynamicUpdateSlice); - - // 2. DynamicSlice tuple index 1: - // Base tensor {1, 0, 0} must remain in default memory (false). - // Prefetched slice output {1, 1} resides in alternate memory (true). - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{while_inst, {1, 0, 0}}, *alias_analysis)); - EXPECT_TRUE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{while_inst, {1, 1}}, *alias_analysis)); - - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{body_param, {1, 0, 0}}, *alias_analysis)); - EXPECT_TRUE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{body_param, {1, 1}}, *alias_analysis)); - - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{root_tuple, {1, 0, 0}}, *alias_analysis)); - EXPECT_TRUE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{root_tuple, {1, 1}}, *alias_analysis)); - - // 3. DynamicUpdateSlice tuple index 2: - // Base tensor {2, 0, 0} and dynamic-update-slice output {2, 1} must remain - // in default memory (false). Update slice {2, 0, 1} resides in alternate - // memory (true). - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{while_inst, {2, 0, 0}}, *alias_analysis)); - EXPECT_TRUE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{while_inst, {2, 0, 1}}, *alias_analysis)); - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{while_inst, {2, 1}}, *alias_analysis)); - - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{body_param, {2, 0, 0}}, *alias_analysis)); - EXPECT_TRUE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{body_param, {2, 0, 1}}, *alias_analysis)); - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{body_param, {2, 1}}, *alias_analysis)); - - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{root_tuple, {2, 0, 0}}, *alias_analysis)); - EXPECT_TRUE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{root_tuple, {2, 0, 1}}, *alias_analysis)); - EXPECT_FALSE(IsAsyncPipelinedWhileAlternateMemoryPosition( - HloPosition{root_tuple, {2, 1}}, *alias_analysis)); -} - -// Test subclass of MsaAlgorithm that simulates an empty candidate list from -// FindChunkCandidates. -class EmptyChunkCandidatesMsaAlgorithm : public MsaAlgorithm { - public: - using MsaAlgorithm::FindBestChunkCandidates; - using MsaAlgorithm::MsaAlgorithm; - - std::vector FindChunkCandidates( - const SlicedBufferInterval& sliced_buffer_interval, - int64_t preferred_offset) const override { - return {}; - } -}; - -// Tests that FindBestChunkCandidates handles empty chunk candidates safely -// without dereferencing past-the-end iterators (preventing ASAN container -// overflows). -TEST_F(MemorySpaceAssignmentTest, FindBestChunkCandidatesEmptyChunkCandidates) { - auto module = CreateNewVerifiedModule(); - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {4}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0}); - ASSERT_OK(module->set_schedule(schedule)); - - AllocationSequence allocations; - Options options = DefaultMemorySpaceOptions(); - ASSERT_OK_AND_ASSIGN(auto alias_analysis, - HloAliasAnalysis::Run(module.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(auto hlo_live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - - EmptyChunkCandidatesMsaAlgorithm algorithm(module.get(), &allocations, - options, *alias_analysis, - &alias_info_, *hlo_live_range); - - AllocationRequest request; - request.end_time = 10; - AliasedOffset preferred_offset{/*offset=*/16}; - - GlobalDecreasingSizeBestFitHeap::BufferInterval buffer_interval; - buffer_interval.buffer = nullptr; - buffer_interval.size = 16; - buffer_interval.start = 0; - buffer_interval.end = 10; - buffer_interval.need_allocation = true; - - using SlicedBufferInterval = - GlobalDecreasingSizeBestFitHeap::SlicedBufferInterval; - auto sliced_interval = - SlicedBufferInterval::CreateMutableInterval(buffer_interval); - - // Without the bounds check in FindBestChunkCandidates, an empty - // chunk_candidates vector results in absl::c_min_element dereferencing an - // end() iterator, triggering an AddressSanitizer container-overflow/crash. - std::vector result = algorithm.FindBestChunkCandidates( - request, &preferred_offset, &sliced_interval); - EXPECT_THAT(result, ::testing::IsEmpty()); -} - -} // namespace - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h deleted file mode 100644 index 7844e734be011e..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h +++ /dev/null @@ -1,634 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_MEMORY_SPACE_ASSIGNMENT_TEST_BASE_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_MEMORY_SPACE_ASSIGNMENT_TEST_BASE_H_ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/check.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/ir/hlo_schedule.h" -#include "xla/hlo/transforms/simplifiers/instruction_hoister.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/buffer_value.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/hlo_buffer.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/buffer_interval_comparator.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.h" -#include "xla/service/memory_space_assignment/options.h" -#include "xla/service/memory_space_assignment/prefetch_interval_picker.h" -#include "xla/service/memory_space_assignment/utils.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/tests/hlo_pjrt_test_base.h" - -namespace xla { -namespace memory_space_assignment { - -constexpr int64_t kPointerSize = 8; -constexpr float kDefaultMemBandwidth = 100; -constexpr float kAlternateMemBandwidth = 1000; -constexpr float kBytesPerSecond = 100; -constexpr float kFlopsPerSecond = 1000; -constexpr float kTranscendentalsPerSecond = 10; - -class TestBufferIntervalComparator : public BufferIntervalComparator { - public: - explicit TestBufferIntervalComparator(MsaBufferIntervalCompare compare_method) - : compare_method_(std::move(compare_method)) {} - - ~TestBufferIntervalComparator() override = default; - - std::string DescribeComparisonCriteria() const override { - return "internal to test"; - } - std::string CriteriaToString( - const MsaBufferInterval& buffer_interval) override { - return "internal to test"; - } - bool LessThan(const MsaBufferInterval& lhs, - const MsaBufferInterval& rhs) override { - return compare_method_(lhs, rhs); - } - - private: - MsaBufferIntervalCompare compare_method_; -}; - -class MemorySpaceAssignmentTestBase : public HloTestBase { - protected: - // We use the following two memory space values to describe the default (slow - // and large) and alternate (fast and small) memory spaces. - const int64_t kDefaultMemorySpace = 0; - const int64_t kAlternateMemorySpace = 1; - const int64_t kPinnedDefaultMemorySpace = 2; - - static HloCostAnalysis::Options DefaultHloCostAnalysisOptions() { - HloCostAnalysis::Options options; - options.set_flops_per_second(kFlopsPerSecond); - options.set_bytes_per_second(kBytesPerSecond); - options.set_transcendentals_per_second(kTranscendentalsPerSecond); - - return options; - } - - Options DefaultMemorySpaceOptions() const { - Options options; - options.max_size_in_bytes = 128; - options.alignment_in_bytes = 8; - // Verify that MSA does not produce overlapping alternate memory - // allocations. This is on by default so that every test is protected - // against chunk-overlap regressions without having to opt in. A test that - // currently trips the verifier must set `verify = false` explicitly, with a - // TODO tracking the underlying bug. - options.verify = true; - options.alternate_memory_space = kAlternateMemorySpace; - options.max_outstanding_prefetches = -1; - options.max_outstanding_evictions = -1; - options.shape_size_fn = HloCostAnalysis::DefaultShapeSize; - - return options; - } - - static CostAnalysisOptions DefaultCostAnalysisOptions() { - CostAnalysisOptions options; - options.default_mem_bandwidth_bytes_per_second = kDefaultMemBandwidth; - options.alternate_mem_read_bandwidth_bytes_per_second = - kAlternateMemBandwidth; - options.alternate_mem_write_bandwidth_bytes_per_second = - kAlternateMemBandwidth; - return options; - } - - static Options UpdateMaxAsyncCopies(Options options, - int64_t max_async_copies) { - options.max_outstanding_prefetches = max_async_copies; - options.max_outstanding_evictions = max_async_copies; - - return options; - } - - // Creates an MsaBufferIntervalCompare function that prioritizes the - // instructions named in prioritized_instruction_names, in the order - // specified. We default to alphabetical instruction name order for the - // remaining instructions. - static MsaBufferIntervalCompare - CreateBufferIntervalCompareFnFromInstructionNames( - std::vector prioritized_instruction_names) { - absl::flat_hash_map instruction_name_to_priority; - // A lower priority value means its higher on the Msa sort list. - for (size_t i = 0; i < prioritized_instruction_names.size(); ++i) { - instruction_name_to_priority[prioritized_instruction_names[i]] = i; - } - return [instruction_name_to_priority = - std::move(instruction_name_to_priority)]( - const MsaBufferInterval& a, const MsaBufferInterval& b) { - auto get_sort_tuple = [&instruction_name_to_priority]( - const MsaBufferInterval& buffer_interval) { - auto it = instruction_name_to_priority.find( - buffer_interval.buffer->defining_instruction()->name()); - if (it != instruction_name_to_priority.end()) { - return std::make_tuple( - it->second, - buffer_interval.buffer->defining_instruction()->name()); - } - return std::make_tuple( - instruction_name_to_priority.size(), - buffer_interval.buffer->defining_instruction()->name()); - }; - - return get_sort_tuple(a) < get_sort_tuple(b); - }; - } - - std::unique_ptr AssignMemorySpaceUsingCostAnalysis( - HloModule* module, - std::optional memory_space_options_override = std::nullopt, - std::optional cost_analysis_options_override = - std::nullopt, - std::optional hlo_cost_options_override = - std::nullopt, - std::optional optional_msa_sort_order_overrides = - std::nullopt) { - HloCostAnalysis::Options hlo_cost_options = DefaultHloCostAnalysisOptions(); - if (hlo_cost_options_override) { - hlo_cost_options = *hlo_cost_options_override; - } - - HloCostAnalysis hlo_cost_analysis(hlo_cost_options); - HloCostAnalysisWithAcceptState hlo_cost_analysis_wrapper(hlo_cost_analysis); - for (HloComputation* computation : module->MakeNonfusionComputations()) { - CHECK_OK(computation->Accept(&hlo_cost_analysis)); - } - auto alias_analysis_or = HloAliasAnalysis::Run(module, &alias_info_); - CHECK_OK(alias_analysis_or.status()); - - Options memory_space_options = DefaultMemorySpaceOptions(); - if (memory_space_options_override) { - memory_space_options = *std::move(memory_space_options_override); - } - CostAnalysisOptions cost_analysis_options = DefaultCostAnalysisOptions(); - if (cost_analysis_options_override) { - cost_analysis_options = *cost_analysis_options_override; - } - OpCostManager op_cost_manager( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - - auto status_or_cost_analysis = - CostAnalysis::Create(op_cost_manager, cost_analysis_options, - &alias_info_, *module, alias_analysis_or->get()); - CHECK_OK(status_or_cost_analysis.status()); - auto cost_analysis = std::move(status_or_cost_analysis.value()); - - memory_space_options.cost_analysis = cost_analysis.get(); - CostAnalysisPrefetchIntervalPicker prefetch_interval_picker( - CostAnalysisPrefetchIntervalPicker( - *cost_analysis, /*min_overlap_to_async_copy_ratio=*/0.8f, - /*preferred_overlap_to_async_copy_ratio=*/1.5, - /*max_overlap_to_mem_size_async_copy_ratio=*/10.0, - /*mem_size_bytes=*/memory_space_options.max_size_in_bytes)); - MsaSortOrderOverrides msa_sort_order_overrides; - if (optional_msa_sort_order_overrides.has_value()) { - msa_sort_order_overrides = optional_msa_sort_order_overrides.value(); - } - MemoryBoundednessBufferIntervalComparator comparator( - *cost_analysis, &cache_, msa_sort_order_overrides); - return AssignMemorySpace( - module, std::move(memory_space_options), - [&comparator](const MsaBufferInterval& lhs, - const MsaBufferInterval& rhs) { - return comparator.LessThan(lhs, rhs); - }, - &prefetch_interval_picker); - } - - std::unique_ptr AssignMemorySpace( - HloModule* module, std::optional options_override = std::nullopt, - int64_t max_prefetch_interval = 10, int64_t min_prefetch_interval = 2) { - InstructionHoister instruction_hoister; - CHECK_OK(instruction_hoister.Run(module).status()); - InstructionCountPrefetchIntervalPicker prefetch_interval_picker( - min_prefetch_interval, max_prefetch_interval); - return AssignMemorySpace(module, std::move(options_override), - /*buffer_interval_compare=*/{}, - &prefetch_interval_picker); - } - - std::unique_ptr AssignMemorySpace( - HloModule* module, std::optional options_override, - std::optional buffer_interval_compare, - PrefetchIntervalPicker* prefetch_interval_picker) { - auto status_or = AssignMemorySpaceAndReturnStatus( - module, std::move(options_override), std::move(buffer_interval_compare), - prefetch_interval_picker); - EXPECT_OK(status_or.status()); - return std::move(status_or.value()); - } - - absl::StatusOr> - AssignMemorySpaceAndReturnStatus( - HloModule* module, std::optional options_override, - std::optional buffer_interval_compare, - PrefetchIntervalPicker* prefetch_interval_picker) { - auto size_fn = [](const BufferValue& buffer) { - return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8); - }; - - auto is_allowed_in_alternate_mem = [](const HloValue& value) { - // Check if the value belongs to the entry computation. - HloInstruction* instruction = value.instruction(); - HloComputation* computation = instruction->parent(); - bool in_entry_computation = - (computation == computation->parent()->entry_computation()); - - return (!in_entry_computation || - instruction->opcode() != HloOpcode::kParameter); - }; - - // Only check parameters in default memory if the original module didn't - // have the parameters in alternate memory. - bool check_parameters_in_default_memory = true; - for (const HloInstruction* parameter : - module->entry_computation()->parameter_instructions()) { - ShapeUtil::ForEachSubshape( - parameter->shape(), - [&](const Shape& subshape, const ShapeIndex& /*index*/) { - if (subshape.has_layout() && - subshape.layout().memory_space() == kAlternateMemorySpace) { - check_parameters_in_default_memory = false; - } - }); - } - - Options options = DefaultMemorySpaceOptions(); - if (options_override) { - options = *std::move(options_override); - } - for (const auto& override : options.msa_tensor_overrides.overrides()) { - if (override.has_pin_in_alternate_memory()) { - check_parameters_in_default_memory = false; - break; - } - } - std::unique_ptr test_comparator; - if (buffer_interval_compare.has_value()) { - test_comparator = std::make_unique( - *buffer_interval_compare); - options.buffer_interval_comparator = test_comparator.get(); - } - options.prefetch_interval_picker = prefetch_interval_picker; - options.size_fn = size_fn; - if (options.is_allowed_in_alternate_mem_fn == nullptr) { - options.is_allowed_in_alternate_mem_fn = is_allowed_in_alternate_mem; - } - - ABSL_ASSIGN_OR_RETURN(auto alias_analysis, - HloAliasAnalysis::Run(module, &alias_info_)); - ABSL_ASSIGN_OR_RETURN(std::unique_ptr hlo_live_range, - HloLiveRange::Run(module->schedule(), *alias_analysis, - module->entry_computation())); - - ABSL_ASSIGN_OR_RETURN( - std::unique_ptr preset_assignments, - MemorySpaceAssignment::Run(module, *hlo_live_range, *alias_analysis, - &alias_info_, options)); - if (check_parameters_in_default_memory) { - CheckParametersInDefaultMemory(module); - } - CheckRootInDefaultMemory(module); - CheckPresetAssignments(preset_assignments.get()); - return preset_assignments; - } - - void CheckPresetAssignments(const PresetAssignments* preset_assignments) { - // Ensure that the exported preset assignments point to layouts in the - // alternate memory. Also ensure that the positions are unique. Note that - // we're using a std::set instead of absl::flat_hash_set because we can make - // use of HloPosition's comparator logic instead of providing a hasher. - std::set positions_in_preset_assignments; - for (auto& position_and_chunk : preset_assignments->chunks()) { - HloPosition position = position_and_chunk.first; - EXPECT_EQ(positions_in_preset_assignments.find(position), - positions_in_preset_assignments.end()); - positions_in_preset_assignments.insert(position); - const Shape& subshape = - ShapeUtil::GetSubshape(position.instruction->shape(), position.index); - EXPECT_EQ(subshape.layout().memory_space(), kAlternateMemorySpace) - << "Exported position is not in alternate mem: " - << position.ToString(); - } - } - - void CheckParametersInDefaultMemory(const HloModule* module) { - // Check that all the entry parameter subshapes are placed in default - // memory. - const HloComputation* entry_computation = module->entry_computation(); - for (const HloInstruction* parameter : - entry_computation->parameter_instructions()) { - ShapeUtil::ForEachSubshape( - parameter->shape(), - [&](const Shape& subshape, const ShapeIndex& /*index*/) { - if (subshape.has_layout()) { - EXPECT_NE(subshape.layout().memory_space(), kAlternateMemorySpace) - << "Parameter not in default memory: " - << parameter->ToString(); - } - }); - } - } - - void CheckRootInDefaultMemory(const HloModule* module) { - const HloInstruction* root = - module->entry_computation()->root_instruction(); - if (root->shape().IsArray()) { - EXPECT_EQ(root->shape().layout().memory_space(), kDefaultMemorySpace); - } - } - - // Returns a std::vector of HloPositions for the given instruction names and - // shape_index. - absl::flat_hash_set GetHloPositions( - const HloModule* module, std::vector instruction_names, - ShapeIndex shape_index = {}) { - absl::flat_hash_set block_prefetched_positions; - for (const auto& instruction_name : instruction_names) { - HloInstruction* param = FindInstruction(module, instruction_name); - EXPECT_NE(param, nullptr); - HloPosition param_position{param, shape_index}; - block_prefetched_positions.insert(param_position); - } - return block_prefetched_positions; - } - - struct CustomCallPrefetchInfo { - std::string prefetched_instruction_name; - std::string prefetch_start_instruction_name; - std::string prefetch_done_instruction_name; - }; - - // Returns a map of HloPositions to CustomCallPrefetchDetails for the given - // custom call prefetch instructions. - absl::flat_hash_map> - GetCustomCallPrefetchDetailsMap( - const HloModule* module, - std::vector custom_call_prefetch_instructions) { - absl::flat_hash_map> - hlo_position_to_custom_call_prefetch_details; - for (const auto& info : custom_call_prefetch_instructions) { - HloInstruction* param = - FindInstruction(module, info.prefetched_instruction_name); - EXPECT_NE(param, nullptr); - HloPosition param_position{param, {}}; - HloInstruction* prefetch_start = - FindInstruction(module, info.prefetch_start_instruction_name); - EXPECT_NE(prefetch_start, nullptr); - HloInstruction* prefetch_done = - FindInstruction(module, info.prefetch_done_instruction_name); - EXPECT_NE(prefetch_done, nullptr); - - CustomCallPrefetchDetails details{/*prefetch_start=*/prefetch_start, - /*prefetch_done=*/prefetch_done, - /*intermediate_instructions=*/ - {prefetch_done->mutable_operand(1), - prefetch_done->mutable_operand(2)}}; - - hlo_position_to_custom_call_prefetch_details[param_position].push_back( - details); - } - return hlo_position_to_custom_call_prefetch_details; - } - - // For each named instruction, the operand at operand_number, should have the - // specified memory space and opcode. - // - // REQUIRES: - // - Each named instruction must have (operand_number+1) operands or more. - // - The operand at operand_number, for each named instruction, must have a - // layout. - // - If operand_opcode is specified, the operand must have the specified - // opcode. - void CheckOperandOpcodeAndMemorySpaceForInstructionNames( - HloModule* module, const std::vector& instruction_names, - int64_t operand_number, int64_t operand_memory_space, - std::optional operand_opcode = std::nullopt) { - for (const std::string& name : instruction_names) { - HloInstruction* use_inst = FindInstruction(module, name); - EXPECT_NE(use_inst, nullptr) << "Instruction not found: " << name; - if (!use_inst) { - continue; - } - const HloInstruction* operand = use_inst->operand(operand_number); - if (operand_opcode.has_value()) { - EXPECT_EQ(operand->opcode(), operand_opcode.value()) - << "Instruction " << operand->name() << " has opcode " - << operand->opcode() << " instead of " << operand_opcode.value(); - } - EXPECT_TRUE(operand->shape().has_layout()) - << "Instruction " << operand->name() << " has no layout."; - if (!operand->shape().has_layout()) { - continue; - } - EXPECT_EQ(operand->shape().layout().memory_space(), operand_memory_space) - << "Instruction " << operand->name() << " has memory space " - << operand->shape().layout().memory_space() << " instead of " - << operand_memory_space; - } - } - - // Checks for every instruction in instruction_names the output memory space - // matches the given memory space. - void CheckMemorySpaceForInstructionNames( - HloModule* module, const std::vector& instruction_names, - int64_t memory_space) { - for (const std::string& name : instruction_names) { - HloInstruction* instruction = FindInstruction(module, name); - EXPECT_NE(instruction, nullptr) << "Instruction not found: " << name; - if (!instruction) { - continue; - } - EXPECT_TRUE(instruction->shape().has_layout()) - << "Instruction " << name << " has no layout."; - if (!instruction->shape().has_layout()) { - continue; - } - EXPECT_EQ(instruction->shape().layout().memory_space(), memory_space) - << "Instruction " << name << " has memory space " - << instruction->shape().layout().memory_space() << " instead of " - << memory_space; - } - } - - struct OutstandingAsyncCopies { - int64_t max_copies; - int64_t max_prefetches; - int64_t max_evictions; - }; - - /*static*/ OutstandingAsyncCopies CountMaximumOutstandingAsyncCopies( - const HloModule& module) const { - OutstandingAsyncCopies copies{0, 0, 0}; - int64_t current_copies = 0; - int64_t current_prefetches = 0; - int64_t current_evictions = 0; - for (HloInstruction* instruction : module.schedule() - .sequence(module.entry_computation()) - .instructions()) { - if (instruction->opcode() == HloOpcode::kCopyStart) { - current_copies++; - if (ShapeUtil::GetSubshape(instruction->shape(), {0}) - .layout() - .memory_space() == kAlternateMemorySpace) { - current_prefetches++; - } else { - current_evictions++; - } - } else if (instruction->opcode() == HloOpcode::kCopyDone) { - current_copies--; - if (instruction->shape().layout().memory_space() == - kAlternateMemorySpace) { - current_prefetches--; - } else { - current_evictions--; - } - } - copies.max_copies = std::max(copies.max_copies, current_copies); - copies.max_prefetches = - std::max(copies.max_prefetches, current_prefetches); - copies.max_prefetches = std::max(copies.max_evictions, current_evictions); - } - return copies; - } - - static int64_t GetAlternateMemoryOffset( - const PresetAssignments& preset_assignments, - const HloInstruction* instruction, const ShapeIndex& index = {}) { - // Returns the offset of the assignment, -1 if it's not in the alternate - // memory. - const HloModule* module = instruction->GetModule(); - AliasInfo alias_info; - auto status_or_alias_analysis = HloAliasAnalysis::Run(module, &alias_info); - CHECK_OK(status_or_alias_analysis.status()); - auto alias_analysis = std::move(status_or_alias_analysis.value()); - const HloBuffer& buffer = - alias_analysis->GetUniqueBufferAt(instruction, index); - for (auto& pos_and_chunk : preset_assignments.chunks()) { - for (auto& value : buffer.values()) { - if (pos_and_chunk.first == value->defining_position()) { - return pos_and_chunk.second.offset; - } - } - } - return -1; - } - - std::unique_ptr CreateEvictAndPrefetchModule() { - HloComputation::Builder builder(TestName()); - Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); - HloInstruction* p0 = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); - HloInstruction* p1 = - builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); - HloInstruction* tanh = builder.AddInstruction( - HloInstruction::CreateUnary(shape, HloOpcode::kTanh, p0)); - // tanh should be placed in the alternate memory since there isn't much - // contention in the beginning. However, tanh has another consumer at the - // end. So it should be kicked out to default memory and prefetched back in. - // The graph below is meant to increase the contention to force - // eviction/prefetch behavior. - HloInstruction* a = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, p0, tanh)); - HloInstruction* b = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* c = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, p0, p1)); - HloInstruction* d = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kSubtract, p0, p1)); - HloInstruction* e = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, b)); - HloInstruction* f = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, c)); - HloInstruction* g = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, a, d)); - HloInstruction* h = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, b, c)); - HloInstruction* i = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, b, d)); - HloInstruction* j = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, c, d)); - HloInstruction* k = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, e, f)); - HloInstruction* l = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, g, h)); - HloInstruction* m = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, i, j)); - HloInstruction* n = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, k, l)); - HloInstruction* o = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, n, m)); - // tanh is being used at the root instruction, and this should be - // prefetched. - HloInstruction* add = builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, o, tanh)); - - auto module = CreateNewVerifiedModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); - - HloSchedule schedule(module.get()); - schedule.set_sequence(computation, {p0, p1, tanh, a, b, c, d, e, f, g, h, i, - j, k, l, m, n, o, add}); - CHECK_OK(module->set_schedule(schedule)); - return module; - } - - CostAnalysis::Cache cache_; - AliasInfo alias_info_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_MEMORY_SPACE_ASSIGNMENT_TEST_BASE_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/options.cc b/third_party/xla/xla/service/memory_space_assignment/options.cc deleted file mode 100644 index 5bc554c0574cef..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/options.cc +++ /dev/null @@ -1,139 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/options.h" - -#include - -#include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" - -namespace xla { -namespace memory_space_assignment { - -std::string PostAllocationTransformationUpdate::ToString() const { - return absl::StrCat("to_be_removed: ", - absl::StrJoin(to_be_removed, ", ", - [](std::string* out, const auto& entry) { - absl::StrAppend(out, entry->name()); - }), - "\n", "update_use_map: ", - absl::StrJoin(update_use_map, ", ", - [](std::string* out, const auto& entry) { - absl::StrAppend( - out, "<", entry.first.ToString(), - " -> ", entry.second.ToString(), ">"); - }), - "\n"); -} - -std::string Options::ToString() const { - return absl::StrJoin( - {absl::StrCat("default_memory_space: ", default_memory_space), - absl::StrCat("alternate_memory_space: ", alternate_memory_space), - absl::StrCat("max_size_in_bytes: ", max_size_in_bytes), - absl::StrCat("alignment_in_bytes: ", alignment_in_bytes), - absl::StrCat("replicated_split_dimension: ", replicated_split_dimension), - absl::StrCat("any_split_dimension: ", any_split_dimension), - absl::StrCat("reduce_scoped_memory_limit: ", reduce_scoped_memory_limit), - absl::StrCat("allocate_reserved_scoped_memory_at_same_offset: ", - allocate_reserved_scoped_memory_at_same_offset), - absl::StrCat("max_outstanding_prefetches: ", max_outstanding_prefetches), - absl::StrCat("max_outstanding_evictions: ", max_outstanding_evictions), - absl::StrCat("while_use_extra_outstanding_prefetch_limit: ", - while_use_extra_outstanding_prefetch_limit), - absl::StrCat("max_retries: ", max_retries), - absl::StrCat("max_repacks: ", max_repacks), - absl::StrCat("repack_after_every_allocation: ", - repack_after_every_allocation), - absl::StrCat("verify: ", verify), - absl::StrCat("enable_cross_program_prefetch: ", - enable_cross_program_prefetch), - absl::StrCat("default_cross_program_prefetch_heuristic: ", - default_cross_program_prefetch_heuristic), - absl::StrCat("enable_cross_program_prefetch_freeing: ", - enable_cross_program_prefetch_freeing), - absl::StrCat("max_cross_program_prefetches: ", - max_cross_program_prefetches), - absl::StrCat("cross_program_prefetch_permissive_mode: ", - cross_program_prefetch_permissive_mode), - absl::StrCat("enable_while_redundant_eviction_elimination: ", - enable_while_redundant_eviction_elimination), - absl::StrCat("use_repeated_instance_for_preferred_prefetch_time: ", - use_repeated_instance_for_preferred_prefetch_time), - absl::StrCat("enforce_prefetch_fifo_order: ", - enforce_prefetch_fifo_order), - absl::StrCat("enable_sync_copy_replacement: ", - enable_sync_copy_replacement), - absl::StrCat("enable_sync_slice_replacement: ", - enable_sync_slice_replacement), - absl::StrCat("enable_sync_custom_fusion_replacement: ", - enable_sync_custom_fusion_replacement), - absl::StrCat("extend_async_copies_limit_for_sync_mem_op_conversion: ", - extend_async_copies_limit_for_sync_mem_op_conversion), - absl::StrCat("inefficient_use_to_copy_ratio: ", - inefficient_use_to_copy_ratio), - absl::StrCat("always_spill_to_default_memory: ", - always_spill_to_default_memory), - absl::StrCat("window_prefetch_mode: ", WindowPrefetchModeToString()), - absl::StrCat("expanded_scoped_alternate_memory_mode: ", - ExpandedScopedAlternateMemoryMode::Value_Name( - expanded_scoped_alternate_memory_mode)), - absl::StrCat("buffer_colorings: ", buffer_colorings.size(), " elements"), - absl::StrCat("post_module_scoped_alternate_memory_size_in_bytes: ", - post_module_scoped_alternate_memory_size_in_bytes), - absl::StrCat("buffer_interval_comparator: ", - buffer_interval_comparator ? "present" : "nullptr"), - absl::StrCat("prefetch_interval_picker: ", - prefetch_interval_picker ? "present" : "nullptr"), - absl::StrCat("cost_analysis: ", cost_analysis ? "present" : "nullptr"), - absl::StrCat("repacker: ", repacker ? "present" : "nullptr"), - absl::StrCat("autotuning_config: ", - autotuning_config.has_value() ? "present" : "nullopt"), - absl::StrCat("preferred_prefetch_overrides: \n", - preferred_prefetch_overrides.DebugString()), - absl::StrCat("msa_tensor_overrides: \n", - msa_tensor_overrides.DebugString()), - absl::StrCat("sliced_prefetch_options: \n", - sliced_prefetch_options.DebugString()), - absl::StrCat("memory_bound_loop_optimizer_options: \n", - memory_bound_loop_optimizer_options.DebugString()), - absl::StrCat("msa_sort_order_overrides: \n", - msa_sort_order_overrides.DebugString())}, - "\n"); -} - -bool Options::IsOpSpanExposureEnabled() const { - return window_prefetch_mode >= WindowPrefetchMode::kAllocationOnly; -} - -bool Options::IsWindowPrefetchingEnabled() const { - return window_prefetch_mode >= WindowPrefetchMode::kPrefetch; -} - -std::string Options::WindowPrefetchModeToString() const { - switch (window_prefetch_mode) { - case WindowPrefetchMode::kPrefetch: - return "prefetch"; - case WindowPrefetchMode::kAllocationOnly: - return "allocation-only"; - case WindowPrefetchMode::kNone: - return "disabled"; - } -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/options.h b/third_party/xla/xla/service/memory_space_assignment/options.h deleted file mode 100644 index 47fb95cc12ef17..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/options.h +++ /dev/null @@ -1,504 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_OPTIONS_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_OPTIONS_H_ - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/layout.h" -#include "xla/service/buffer_value.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation_value.h" -#include "xla/service/memory_space_assignment/buffer_interval_comparator.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/memory_space_assignment/prefetch_interval_picker.h" -#include "xla/service/memory_space_assignment/repacking.h" -#include "xla/service/memory_space_assignment/slice.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -using IsAllowedInAlternateMemoryFunction = std::function; -using IsUseAllowedInAlternateMemoryFunction = - std::function; -using IsPositionAllowedInAlternateMemoryFunction = - std::function; -using ReservedScopedMemoryFunction = std::function>& /*operands_in_alternate_memory*/, - const absl::flat_hash_set& /*outputs_in_alternate_memory*/)>; -using PositionRequiresContiguousAllocationFunction = - std::function; -using WindowPrefetchNotifyOperandAppendedFunction = - std::function; -using IsAsyncSliceImplementedFunction = - std::function; -using InitSplitTreeFn = std::function( - const HloInstruction*, - absl::flat_hash_map>*)>; -using DetermineSplitDimensionFunction = - std::function( - const HloValue&, - absl::flat_hash_map>*)>; -using BitcastSplitFn = std::function( - const HloInstruction* instruction, int64_t split_dim)>; -using ShapeSizeFn = std::function; -using AsyncInstructionBwAdjustmentFactorFn = - std::function(const HloInstruction*)>; -using IsWindowPrefetchableInstructionFunction = - std::function; -using HloPositionOrUse = std::variant; -using OpSpanSizeFn = std::function; - -// MSA allows for custom post-allocation transformations. When a post-allocation -// transformation is performed on an instruction, this result is returned. It -// tells MSA: -// 1. A list of instructions that MSA should delete. -// 2. A list of HloUses that the transformation replaced. -// -// This information is then processed via -// FixAllocationSequenceAfterPostAllocationTransformation call. -struct PostAllocationTransformationUpdate { - std::vector to_be_removed; - absl::flat_hash_map update_use_map; - - std::string ToString() const; -}; - -// The different modes for window prefetch. kAllocationOnly allocates window -// buffers separately from scoped alternate memory. kPrefetch not only allocates -// independent window buffers, but it also prefetches their initial content. -enum class WindowPrefetchMode { - kNone, - kAllocationOnly, - kPrefetch, -}; - -// A struct to specify the memory space coloring of a buffer position or use. -struct BufferColoring { - HloPositionOrUse buffer_position_or_use; // Buffer position or use to color. - int64_t memory_space; // How to color the buffer. -}; - -// A struct to hold the details of a custom call prefetch. -struct CustomCallPrefetchDetails { - // Async custom call prefetch start instruction. - HloInstruction* prefetch_start; - // Async custom call prefetch done instruction. - HloInstruction* prefetch_done; - // Intermediate instructions associated with the prefetch like - // get-tuple-element etc. - std::vector intermediate_instructions; -}; - -// The different options to be passed to the Run() API. -struct Options { - std::string ToString() const; - - // Returns true if op span exposure is enabled. - bool IsOpSpanExposureEnabled() const; - - // Returns true if window prefetching is enabled. - bool IsWindowPrefetchingEnabled() const; - - std::string WindowPrefetchModeToString() const; - - // The execution threads that memory space assignment operates on. This is - // used to determine which flattened instructions should be included/excluded - // from cost analysis. - absl::flat_hash_set execution_threads = { - HloInstruction::kMainExecutionThread}; - - // The backend-specific integer value that describes the default memory. - int64_t default_memory_space = 0; - - // Backend-specific integer value that describes the alternate memory. - int64_t alternate_memory_space = 0; - - // Color of "view" values. Views are zero-copy, possibly sub-region aliases - // into another allocation. We identify them by their color. MSA extends the - // pointed-to buffer's allocation to account for the live range of views to - // it. Pair with an is_allowed_in_alternate_mem_fn that rejects view colored - // values. std::nullopt disables views. - std::optional dus_view_color; - - // Maximum size of the alternate memory space. - int64_t max_size_in_bytes = 0; - - // Memory alignment of the alternate memory space. - int64_t alignment_in_bytes = 1; - - // If provided, we sort the buffers using this comparator. Otherwise, we use - // GlobalDecreasingSizeBestFitHeap::kSpatial. - BufferIntervalComparator* buffer_interval_comparator = nullptr; - - // This object determines how early and how late prefetches can occur. - PrefetchIntervalPicker* prefetch_interval_picker = nullptr; - - // This object is used to determine the benefit of a particular allocation. - CostAnalysis* cost_analysis = nullptr; - - // Size function for buffer values. - BufferValue::SizeFunction size_fn; - - ShapeSizeFn shape_size_fn; - - std::function get_equivalent_s8_shape_fn; - - // This function can be used to prevent certain HloValues (e.g., based on - // the opcode) to be placed on the alternate memory. - IsAllowedInAlternateMemoryFunction is_allowed_in_alternate_mem_fn; - - // This function can be used to prevent certain HloUses (e.g., based on - // the opcode) to be placed on the alternate memory. - IsUseAllowedInAlternateMemoryFunction is_use_allowed_in_alternate_mem_fn = - [](const HloUse&) { return true; }; - - // Specifies if the given position is allowed in the alternate memory. - IsPositionAllowedInAlternateMemoryFunction - is_position_allowed_in_alternate_mem_fn = - [](const HloPosition&) { return true; }; - - // This function returns the amount of scoped memory in bytes that should be - // reserved during the execution of this instruction. Note that the - // `operands_in_alternate_memory` also includes the window prefetched - // operands. - ReservedScopedMemoryFunction reserved_scoped_memory_fn = - [](const HloInstruction*, - const absl::flat_hash_set< - std::pair>& /*operands_in_alternate_memory*/, - const absl::flat_hash_set< - ShapeIndex>& /*outputs_in_alternate_memory*/) { return 0; }; - - PositionRequiresContiguousAllocationFunction - position_requires_contiguous_allocation_fn = - [](const HloPosition&) { return false; }; - - // This function is used to determine the size of the span buffer for a given - // operand. The size should be the total size required by the operand. If - // pipelining is disabled, this is one iteration worth of data; if pipelining - // is enabled and double buffering is used, this is two iterations worth of - // data. - // - // `hlo_with_memory_spaces` is a clone of `original_hlo` but with the memory - // space assignment propagated within its computation. It is used to determine - // the operand span size of the operand. `operand_index` is the index of the - // operand which is being considered. - OpSpanSizeFn op_span_size_fn = [](HloInstruction* original_hlo, - HloInstruction* hlo_with_memory_spaces, - int64_t operand_index) { return 0; }; - - // This function is called to notify that an operand has been appended as a - // window prefetch buffer. - WindowPrefetchNotifyOperandAppendedFunction notify_operand_appended_fn = - [](HloInstruction*, int64_t, int64_t) {}; - - // This function can be used to check if an equivalent asynchronous slice - // lowering is implemented for a given synchronous slice instruction. - IsAsyncSliceImplementedFunction is_async_slice_implemented_fn = - [](const HloInstruction*) { return false; }; - - // Should only be used for testing purposes. This function allows us to - // modify the AllocationResult after the AllocationRequest has been processed - // by AllocateSegment(). - std::function - allocation_result_modifier_testing_fn = nullptr; - - // Should only be used for testing purposes. This function allows us to - // modify the AllocationRequest before the AllocationRequest is passed to - // AllocateSegment(). - std::function - allocation_request_modifier_testing_fn = nullptr; - - // This function chooses a dimension to split the given HloValue on. Splitting - // will be disabled if this function is not provided. - DetermineSplitDimensionFunction determine_split_dimension_fn = nullptr; - - // This function sets up a split tree, based on an instruction's shape, with - // kAny as the default. Splitting will be disabled if this function is not - // provided. - InitSplitTreeFn init_split_tree_fn = nullptr; - - // Determines the appropriate output split for a bitcast given an input split. - // Splitting will be disabled if this function is not provided. - BitcastSplitFn bitcast_split_fn = nullptr; - - // Dimension number indicating no split is present. - int64_t replicated_split_dimension = -1; - - // Dimension number indicating any split is allowable. - int64_t any_split_dimension = -2; - - // Applies post-allocation transformations to the given instruction. This - // function is called after the allocations are found in the MsaAlgorithm. It - // is called on each instruction I that meets the following conditions: - // 1. I is called from a non-fusion computation - // 2. I's operands are not in alternate memory - // 3. I is not successfully converted to async instruction. - // 4. I's operands don't have in-place users, e.g., a dynamic-update-slice. - // - // The transformation function is allowed to do the following: - // 1. Mark instructions for removal. - // 2. Modify existing instructions. - // - // This transformation is NOT allowed to: - // 1. Directly remove instructions (or nullify them). - // 2. Add new instructions. - // - // Note that it is up to the transformation function to ensure that the - // changes to the module preserves the semantics of the original program. - std::function( - HloInstruction*)> - post_allocation_transformation_fn; - - // If true, we will try to reduce scoped allocation buffer size for all - // instructions if their operand/output has been allocated in alternate - // memory. - bool reduce_scoped_memory_limit = false; - - // If true, we allocate the reserved scoped memory at the same offset. This - // is useful to enable more deduplication between HLOs that have reserved - // scoped memories, but may result in less efficient memory packing. - bool allocate_reserved_scoped_memory_at_same_offset = true; - - // Specifies the upper bound for number of outstanding prefetches and - // evictions, -1 for unlimited. - int64_t max_outstanding_prefetches = -1; - int64_t max_outstanding_evictions = -1; - - // Extra outstanding prefetch limit for while uses (in addition to - // max_outstanding_prefetches). - int64_t while_use_extra_outstanding_prefetch_limit = 0; - - // Specifies the maximum number of retries that will be performed for each - // value in case prefetching failed due to running out of asynchronous - // copies or asynchronous copy resource. - int64_t max_retries = 1; - - // The maximum number of repacks that we are willing to perform in case we - // can't allocate a buffer due to running out of memory. If this value is - // greater than 0, repacker must be non-nullptr. - int64_t max_repacks = 0; - - // The repacking algorithm to reduce fragmentation. Must be non-null if - // max_repacks is greater than 0. - MemorySpaceAssignmentRepacker* repacker = nullptr; - - // This is only useful for testing, repack after every allocation. - bool repack_after_every_allocation = false; - - // If true, verifies the memory space assignment against overlapping - // buffers. - bool verify = false; - - // If not nullptr, this function is called to dump debugging information. - // The first argument is appended to the file name and the second argument - // is the contents of the file. - std::function dump_fn = nullptr; - - // Enable prefetching buffers into preferred memory across program - // boundaries - bool enable_cross_program_prefetch = true; - - // If true, use buffer_interval_compare to determine which buffers to - // prefetch across program boundaries. - bool default_cross_program_prefetch_heuristic = false; - - // Enable cross-program prefetch freeing optimization where the - // cross-program-prefetched buffer can be reused. - bool enable_cross_program_prefetch_freeing = true; - - // The maximum number of cross program prefetches. - // TODO(tjablin): Use a heuristic to determine this automatically. - int max_cross_program_prefetches = 1; - - // If false, we assume tensors that we couldn't explicitly determine to be - // activations are activations. If true, we assume these aren't activations, - // so they may be cross-program-prefetch candidates. - bool cross_program_prefetch_permissive_mode = false; - - // Enable redundant eviction optimization in/around while loops. If enabled, - // this optimization would keep a copy of the buffer in the default memory in - // addition to alternate memory to eliminate redundant evictions. - bool enable_while_redundant_eviction_elimination = true; - - // An optional memory space assignment autotuning config, which is used - // to sort allocated buffers. - std::optional> autotuning_config = std::nullopt; - - // If true, uses the earlier instance of the same instruction to use as - // preferred prefetch start time. - bool use_repeated_instance_for_preferred_prefetch_time = false; - - // If true, enforces the FIFO order for prefetches. - bool enforce_prefetch_fifo_order = false; - - // If true, tries to replace synchronous copy instructions with asynchronous - // ones. If it fails to replace the copy, it keeps the sync version. - bool enable_sync_copy_replacement = false; - - // If true, tries to replace synchronous slice instructions with asynchronous - // ones. If it fails to replace the slice, it keeps the sync version. - bool enable_sync_slice_replacement = false; - - // Used in block prefetching mode. If true, try to asyncify DMA like custom - // fusion instructions, like cross-buffer slice fusions. - bool enable_sync_custom_fusion_replacement = false; - - // If non-zero, this is the number of extra outstanding async copies that we - // allow for each sync mem op that is converted to an async mem op. - int extend_async_copies_limit_for_sync_mem_op_conversion = 0; - - // The ratio of use bytes to copy bytes for a given allocation site below - // which we consider the site to be inefficient. A value of 0 would treat all - // sites as efficient and a value of 1 would require the amount of bytes used - // at the site to be at least as much as the async copy bytes. There are two - // factors that determine the copy and use bytes: - // - Some uses don't actually access the entire tensor, e.g. in - // dynamic-update-slice. - // - copy_bytes may be larger than the size of the tensor as well. An - // example is a tensor may be prefetched, used, and then evicted. In that - // case copy_bytes would be twice the size of the tensor. - float inefficient_use_to_copy_ratio = 0.0; - - // This is mostly used for testing, it allows a test case to inject its own - // logic for MsaAlgorithm::GetInefficientAllocationSites. - std::function>( - absl::Span)> - get_inefficient_allocation_sites_fn = nullptr; - - // Config to filter prefetches and update preferred prefetch times for the - // filtered prefetches. - PreferredPrefetchOverrides preferred_prefetch_overrides; - - // Unified overrides for tensor placement, pinning, retention, and prefetch - // timing. - MsaTensorOverrides msa_tensor_overrides; - - // Options for slicing prefetches into smaller asynchronously copied pieces. - SlicedPrefetchOptions sliced_prefetch_options; - - // Options for the memory-bound loop optimizer feature. - MemoryBoundLoopOptimizerOptions memory_bound_loop_optimizer_options; - - SliceProposalFunction propose_slice_fn = [](const Shape&, - const SlicedPrefetchOptions&) - -> absl::StatusOr { - return UnimplementedStrCat("Generation of SliceProposals unimplemented"); - }; - - // Option to always spill buffers from alternate memory to default memory - // and prefetching back to alternate memory(if needed) just in time for use. - bool always_spill_to_default_memory = false; - - // Max number of window prefetch operands allowed. - int64_t window_prefetch_max_operands = 1024; - - // Min span size for window prefetch operands allowed. - int64_t window_prefetch_min_span_size = 4096; - - // This function is called to determine whether an instruction is eligible - // for window prefetching. If not provided (nullptr), the default logic is - // used: output fusions and loop fusions not on sparsecore are eligible. - IsWindowPrefetchableInstructionFunction - is_window_prefetchable_instruction_fn = - [](const HloInstruction* instruction) { - return instruction->IsOutputFusion() || instruction->IsLoopFusion(); - }; - - // The mode to use for window prefetching. - WindowPrefetchMode window_prefetch_mode = WindowPrefetchMode::kNone; - - MsaSortOrderOverrides msa_sort_order_overrides; - - // If true, allocates alternate memory colored buffers after pre-colored - // buffers but before other buffers. - bool allocate_colored_buffers_early = true; - - // A mode that enables expanding scoped alternate memory allocations to the - // largest contiguous open space available. - ExpandedScopedAlternateMemoryMode::Value - expanded_scoped_alternate_memory_mode = - ExpandedScopedAlternateMemoryMode::DISABLED; - - std::vector buffer_colorings; - - // If set, this is the size of scoped alternate memory that we require MSA to - // allocate for post-module operations. - uint64_t post_module_scoped_alternate_memory_size_in_bytes = 0; - - // This is the maximum number of concurrent block prefetches allowed. - int64_t max_outstanding_block_prefetches = 0; - - // This is the size of alternate memory that available for block prefetches. - uint64_t reserved_bytes_for_block_prefetches = 0; - - // List of hlo positions for block prefetches. - absl::flat_hash_set block_prefetched_positions; - - // Determines the bandwidth adjustment factor for an async start instruction. - // The available bandwidth for instructions between this and the async done - // instruction will be multiplied by the factor returned by this function. A - // factor of 1.0 means that the full bandwidth is available. A factor of 0.5 - // means that only half the bandwidth is available. - AsyncInstructionBwAdjustmentFactorFn - async_instruction_bw_adjustment_factor_fn = - [](const HloInstruction*) { return std::nullopt; }; - - // One HloPosition can have multiple custom call prefetches associated with - // it. For every custom-call prefetched HloPosition, this map stores the - // details of all the custom-call prefetches associated with it. This is used - // to schedule the custom-call prefetches as part of the memory space - // assignment algorithm. - absl::flat_hash_map> - hlo_position_to_custom_call_prefetch_details; - - // Used in block prefetching mode. Returns the operand index of the source - // buffer (source of a DMA) for a custom fusion instruction that can be - // asyncified. Currently this matches cross-buffer slice fusions which can - // be lowered to asynchronous DMAs. - std::function(const HloInstruction*)> - custom_fusion_block_prefetch_operand_index_fn = - [](const HloInstruction*) { return std::nullopt; }; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_OPTIONS_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc deleted file mode 100644 index 4a38f949752eed..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc +++ /dev/null @@ -1,587 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/prefetch_interval_picker.h" - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/strings/str_cat.h" -#include "absl/types/span.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "xla/util.h" -#include "xla/xla_data.pb.h" -#include "tsl/platform/logging.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -// Each time we retry compilation, increase the preferred eviction end time by -// this amount multiplied by preferred overlap to async copy ratio. -const float kEvictionRetryMultiplier = 2.0; - -// The number of decreasing intervals for CostAnalysisPrefetchIntervalPicker to -// return when it runs out of increasing intervals. Increasing this number may -// hurt compilation time. -const int kNumExploredDecreasingIntervals = 100; - -} // namespace - -bool InstructionCountPrefetchIntervalPicker::CanAllocateInAlternateMemoryNoCopy( - int64_t shape_size, int64_t start_time, int64_t end_time) const { - return end_time - start_time <= max_overlap_count_; -} - -int64_t InstructionCountPrefetchIntervalPicker::PreferredEvictionEndTime( - const Shape& shape, int64_t start_time, int64_t latest_end_time) const { - return std::min(start_time + min_overlap_count_, latest_end_time); -} - -int64_t InstructionCountPrefetchIntervalPicker::LatestPrefetchStartTime( - int64_t shape_size, int64_t start_time, int64_t end_time, - const HloUse* use) const { - return end_time - min_overlap_count_; -} - -int64_t InstructionCountPrefetchIntervalPicker::PreferredPrefetchStartTime( - const Shape& shape, int64_t earliest_prefetch_start_time, - int64_t latest_prefetch_start_time, int64_t prefetch_end_time) const { - return std::max(earliest_prefetch_start_time, - prefetch_end_time - max_overlap_count_); -} - -int64_t InstructionCountPrefetchIntervalPicker::EstimatedPrefetchEndTime( - const Shape& shape, int64_t start_time, int64_t end_time) const { - // For testing, assume the end time is the estimated prefetch end time. - return end_time; -} - -float InstructionCountPrefetchIntervalPicker::GetLogicalIntervalElapsed( - int64_t start_time, int64_t end_time) const { - // For testing, just assume every HLO takes 1 second. - return static_cast(end_time - start_time - 1); -} - -void InstructionCountPrefetchIntervalPicker::Begin( - const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time, bool strict_timing) { - end_time_ = end_time; - strict_timing_ = strict_timing; - const Shape& shape = ShapeUtil::GetSubshape( - use.instruction->operand(use.operand_number)->shape(), use.operand_index); - if (preferred_time) { - current_prefetch_time_ = *preferred_time; - if (strict_timing) { - strict_prefetch_time_ = *preferred_time; - } - } else { - current_prefetch_time_ = - PreferredPrefetchStartTime(shape, start_time, end_time, end_time); - if (strict_timing) { - strict_prefetch_time_ = current_prefetch_time_; - } - } -} - -int64_t InstructionCountPrefetchIntervalPicker::Next() { - CHECK(!Done()) << "Prefetch interval picker's Next() is called even though " - "Done() is false"; - return current_prefetch_time_++; -} - -bool InstructionCountPrefetchIntervalPicker::Done() const { - if (strict_timing_) { - return current_prefetch_time_ > strict_prefetch_time_; - } - return end_time_ - current_prefetch_time_ <= min_overlap_count_; -} - -int64_t InstructionCountPrefetchIntervalPicker::latest_time() const { - if (strict_timing_) { - return strict_prefetch_time_; - } - return end_time_ - min_overlap_count_ - 1; -} - -std::string InstructionCountPrefetchIntervalPicker::ToDebugString() const { - return absl::StrCat("Overlapped HLOs = ", end_time_ - current_prefetch_time_); -} - -std::string InstructionCountPrefetchIntervalPicker::ToNoCopyDebugString( - const Shape& shape, int64_t start_time, int64_t end_time) const { - return absl::StrCat("Overlapped HLOs = ", end_time - start_time); -} - -CostAnalysisPrefetchIntervalPicker::CostAnalysisPrefetchIntervalPicker( - const CostAnalysis& cost_analysis, float min_overlap_to_async_copy_ratio, - float preferred_overlap_to_async_copy_ratio, - float max_overlap_to_mem_size_async_copy_ratio, int64_t mem_size_bytes, - std::optional size_override) - : while_nest_level_( - cost_analysis.hlo_live_range().instruction_schedule().size() + 1, 0), - computation_nest_level_( - cost_analysis.hlo_live_range().instruction_schedule().size() + 1, 0), - cost_analysis_(cost_analysis), - min_overlap_to_async_copy_ratio_(min_overlap_to_async_copy_ratio), - preferred_overlap_to_async_copy_ratio_( - preferred_overlap_to_async_copy_ratio), - size_override_(size_override) { - max_async_copy_elapsed_ = - cost_analysis_.GetAsyncCopyElapsed(cost_analysis_.GetShapeSizeBytes( - ShapeUtil::MakeShape(S32, {mem_size_bytes / 4}))) * - max_overlap_to_mem_size_async_copy_ratio; - instruction_schedule_ = - &cost_analysis_.hlo_live_range().instruction_schedule(); - - // Create a vector of elapsed times and while nesting levels of HLO - // instructions. The elapsed times are multiplied by - // pow(while_execution_count, nest_level) to account for executing the HLOs - // multiple times in while loops. - std::vector instructions_elapsed_time( - instruction_schedule_->size() + 1, 0.0); - int max_while_nest_level = 0; - for (const auto& instruction_and_logical_time : *instruction_schedule_) { - // To avoid double counting, don't include the elapsed time of while and - // conditional HLOs. - const HloInstruction* instruction = instruction_and_logical_time.first; - int64_t logical_time = instruction_and_logical_time.second; - if (logical_time >= instructions_elapsed_time.size()) { - instructions_elapsed_time.resize(logical_time + 1, 0.0); - while_nest_level_.resize(logical_time + 1, 0); - } - int while_nest_level = cost_analysis_.CalculateComputationNestLevel( - instruction_and_logical_time.first, /*while_only=*/true); - while_nest_level_[logical_time] = while_nest_level; - max_while_nest_level = std::max(max_while_nest_level, while_nest_level); - int computation_nest_level = cost_analysis_.CalculateComputationNestLevel( - instruction_and_logical_time.first, /*while_only=*/false); - computation_nest_level_[logical_time] = computation_nest_level; - if (instruction->opcode() == HloOpcode::kWhile || - instruction->opcode() == HloOpcode::kConditional) { - continue; - } - float elapsed_time = cost_analysis_.GetInstructionElapsed( - *instruction_and_logical_time.first); - instructions_elapsed_time[logical_time] = - elapsed_time * cost_analysis_.GetWhileNestMultiplier(while_nest_level); - } - // As an optimization, create a cumulative sum vector of elapsed time. - float cumsum = 0.0; - elapsed_time_cumsum_.reserve(instructions_elapsed_time.size()); - for (float elapsed_time : instructions_elapsed_time) { - cumsum += elapsed_time; - elapsed_time_cumsum_.push_back(cumsum); - } - // To be able to accurately determine the minimum nest level between a start - // time and an end time efficiently, populate a data structure that stores the - // closest 'smaller' nest level change index. - const int64_t size = instructions_elapsed_time.size(); - CHECK_EQ(size, while_nest_level_.size()); - std::vector most_recent_by_level(while_nest_level_.size(), -1); - int prev_nest_level = 0; - int change_idx = -1; - while_nest_level_change_.reserve(size); - for (int i = 0; i < size; ++i) { - int nest_level = while_nest_level_[i]; - if (nest_level != prev_nest_level) { - prev_nest_level = nest_level; - // Compute last change index by choosing the most recent instruction index - // with smaller nesting level. Note that it may happen that even though - // there were few different regions with other nest levels before, all of - // then are same or bigger than this one, in which case we'll end up with - // -1, e.g. if you got nest level 0 no need checking anything else. - change_idx = -1; - for (int smaller_level = 0; smaller_level < nest_level; smaller_level++) { - change_idx = std::max(change_idx, most_recent_by_level[smaller_level]); - } - } - most_recent_by_level[nest_level] = i; - while_nest_level_change_.push_back(change_idx); - } - for (int i = 0; i <= max_while_nest_level; ++i) { - while_execution_counts_.push_back(cost_analysis_.GetWhileNestMultiplier(i)); - } -} - -float CostAnalysisPrefetchIntervalPicker::GetMaxElapsedInAlternateMemory( - float async_copy_elapsed) const { - return max_async_copy_elapsed_; -} - -bool CostAnalysisPrefetchIntervalPicker::CanAllocateInAlternateMemoryNoCopy( - int64_t shape_size, int64_t start_time, int64_t end_time) const { - // Even though this method returns if we allow the buffer in alternate memory - // _without_ asynchronous copies, calculate how long it would have taken to - // copy it and compare it to the elapsed time in the logical interval. - int64_t effective_size = size_override_ ? *size_override_ : shape_size; - float async_copy_elapsed = static_cast(effective_size) / - cost_analysis_.DefaultMemBandwidthBytesPerSecond( - /*use_scaling_factor=*/true); - float logical_interval_elapsed = - GetLogicalIntervalElapsed(start_time, end_time); - return GetMaxElapsedInAlternateMemory(async_copy_elapsed) > - logical_interval_elapsed; -} - -int64_t CostAnalysisPrefetchIntervalPicker::PreferredEvictionEndTime( - const Shape& shape, int64_t start_time, int64_t latest_end_time) const { - float async_copy_elapsed = cost_analysis_.GetAsyncCopyElapsed( - size_override_ ? *size_override_ - : cost_analysis_.GetShapeSizeBytes(shape)); - int64_t end_time; - for (end_time = start_time + 1; end_time <= latest_end_time; ++end_time) { - float logical_interval_elapsed = - GetLogicalIntervalElapsed(start_time, end_time); - if (logical_interval_elapsed >= - (1 + kEvictionRetryMultiplier * retry_number_) * - preferred_overlap_to_async_copy_ratio_ * async_copy_elapsed) { - break; - } - } - return end_time; -} - -int64_t CostAnalysisPrefetchIntervalPicker::LatestPrefetchStartTime( - int64_t shape_size, int64_t start_time, int64_t end_time, - const HloUse* use) const { - // Find the earliest time that satisfies max_overlap_to_async_copy_ratio_. - float async_copy_elapsed = cost_analysis_.GetAsyncCopyElapsed( - size_override_ ? *size_override_ : shape_size); - // If there is a use, estimate the time we would save by having this op in - // alternate memory. - float inst_elapsed_reduction = 0.0f; - if (use) { - float elapsed_time = - cost_analysis_.GetInstructionElapsed(*use->instruction); - float elapsed_time_in_alternate_mem = - cost_analysis_.GetInstructionElapsedInAlternateMemory( - *use->instruction, - /*operands_in_alternate_mem=*/ - {std::make_pair(use->operand_number, use->operand_index)}, - /*outputs_in_alternate_mem=*/{}); - inst_elapsed_reduction = elapsed_time - elapsed_time_in_alternate_mem; - } - int end_nest_level = computation_nest_level_[end_time]; - - // Find the latest time we're allowed to start prefetching. - float min_interval = min_overlap_to_async_copy_ratio_ * async_copy_elapsed; - int latest_prefetch_time; - for (latest_prefetch_time = end_time - 1; - latest_prefetch_time >= start_time && - (computation_nest_level_[latest_prefetch_time] != end_nest_level || - min_interval > - GetLogicalIntervalElapsed(latest_prefetch_time, end_time) + - inst_elapsed_reduction); - --latest_prefetch_time) { - } - - return latest_prefetch_time; -} - -int64_t CostAnalysisPrefetchIntervalPicker::PreferredPrefetchStartTime( - const Shape& shape, int64_t earliest_prefetch_start_time, - int64_t latest_prefetch_start_time, int64_t prefetch_end_time) const { - // Between the earliest and latest prefetch interval, find the interval - // closest to the preferred interval and start iterating from there. - float async_copy_elapsed = cost_analysis_.GetAsyncCopyElapsed( - size_override_ ? *size_override_ - : cost_analysis_.GetShapeSizeBytes(shape)); - int64_t preferred_prefetch_start_time = earliest_prefetch_start_time; - float preferred_interval = - preferred_overlap_to_async_copy_ratio_ * async_copy_elapsed; - float best_interval = GetLogicalIntervalElapsed(earliest_prefetch_start_time, - prefetch_end_time); - int end_nest_level = computation_nest_level_[prefetch_end_time]; - for (int64_t prefetch_start_time = earliest_prefetch_start_time + 1; - prefetch_start_time <= latest_prefetch_start_time; - ++prefetch_start_time) { - float interval = - GetLogicalIntervalElapsed(prefetch_start_time, prefetch_end_time); - if (computation_nest_level_[prefetch_start_time] == end_nest_level && - std::abs(preferred_interval - interval) < - std::abs(preferred_interval - best_interval)) { - best_interval = interval; - preferred_prefetch_start_time = prefetch_start_time; - } - } - return preferred_prefetch_start_time; -} - -int64_t CostAnalysisPrefetchIntervalPicker::LatestPrefetchEndTime( - int64_t original_prefetch_end_time, - int64_t proposed_prefetch_end_time) const { - // Iterate towards the beginning until we find a suitable end time that is the - // same while nest level as the original prefetch end time. - int64_t original_nest_level = - computation_nest_level_[original_prefetch_end_time]; - int64_t new_prefetch_end_time; - for (new_prefetch_end_time = proposed_prefetch_end_time; - computation_nest_level_[new_prefetch_end_time] != original_nest_level; - --new_prefetch_end_time) { - } - return new_prefetch_end_time; -} - -int64_t CostAnalysisPrefetchIntervalPicker::EstimatedPrefetchEndTime( - const Shape& shape, int64_t start_time, int64_t end_time) const { - float async_copy_elapsed = cost_analysis_.GetAsyncCopyElapsed( - size_override_ ? *size_override_ - : cost_analysis_.GetShapeSizeBytes(shape)); - int64_t estimated_end_time; - for (estimated_end_time = start_time + 1; estimated_end_time < end_time; - ++estimated_end_time) { - float interval = GetLogicalIntervalElapsed(start_time, estimated_end_time); - if (interval >= async_copy_elapsed) { - break; - } - } - return estimated_end_time; -} - -void CostAnalysisPrefetchIntervalPicker::Begin( - const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time, bool strict_timing) { - const Shape& shape = ShapeUtil::GetSubshape( - use.instruction->operand(use.operand_number)->shape(), use.operand_index); - int64_t shape_size = size_override_ ? *size_override_ - : cost_analysis_.GetShapeSizeBytes(shape); - // Find the earliest time that satisfies max_overlap_to_async_copy_ratio_. - async_copy_elapsed_ = cost_analysis_.GetAsyncCopyElapsed(shape_size); - // Estimate the time we would save by having this op in alternate memory. - float elapsed_time = cost_analysis_.GetInstructionElapsed(*use.instruction); - float elapsed_time_in_alternate_mem = - cost_analysis_.GetInstructionElapsedInAlternateMemory( - *use.instruction, /*operands_in_alternate_mem=*/ - {std::make_pair(use.operand_number, use.operand_index)}, - /*outputs_in_alternate_mem=*/{}); - inst_elapsed_reduction_ = elapsed_time - elapsed_time_in_alternate_mem; - end_logical_time_ = end_time; - int end_nest_level = computation_nest_level_[end_logical_time_]; - - if (strict_timing && preferred_time.has_value()) { - int64_t target_time = *preferred_time; - earliest_prefetch_time_ = target_time; - latest_prefetch_time_ = target_time; - increasing_prefetch_time_iterator_ = target_time; - decreasing_prefetch_time_iterator_ = target_time; - using_increasing_prefetch_time_iterator_ = true; - Next(); - return; - } - - // Find the latest time we're allowed to start prefetching. - float min_interval = min_overlap_to_async_copy_ratio_ * async_copy_elapsed_; - latest_prefetch_time_ = - LatestPrefetchStartTime(shape_size, start_time, end_time, &use); - - // Find the earliest time we're allowed to start prefetching. - float max_interval = GetMaxElapsedInAlternateMemory(async_copy_elapsed_); - for (earliest_prefetch_time_ = start_time; - earliest_prefetch_time_ < latest_prefetch_time_ && - (computation_nest_level_[earliest_prefetch_time_] != end_nest_level || - max_interval < GetLogicalIntervalElapsed(earliest_prefetch_time_, - end_logical_time_)); - ++earliest_prefetch_time_) { - } - if (earliest_prefetch_time_ > latest_prefetch_time_) { - // There is no available prefetch interval for the given start and end - // times. Set the iterators accordingly to ensure Done() returns true. - increasing_prefetch_time_iterator_ = earliest_prefetch_time_; - decreasing_prefetch_time_iterator_ = latest_prefetch_time_; - CHECK(Done()); - return; - } - - int64_t starting_prefetch_time; - if (preferred_time && *preferred_time <= latest_prefetch_time_) { - starting_prefetch_time = *preferred_time; - } else { - starting_prefetch_time = - PreferredPrefetchStartTime(shape, earliest_prefetch_time_, - latest_prefetch_time_, end_logical_time_); - } - float preferred_interval = - preferred_overlap_to_async_copy_ratio_ * async_copy_elapsed_; - VLOG(4) << "Interval min/max/preferred = " << min_interval << " " - << max_interval << " " << preferred_interval - << " prefetch time earliest/latest/starting = " - << earliest_prefetch_time_ << " " << latest_prefetch_time_ << " " - << starting_prefetch_time; - - increasing_prefetch_time_iterator_ = starting_prefetch_time; - decreasing_prefetch_time_iterator_ = starting_prefetch_time; - using_increasing_prefetch_time_iterator_ = true; - // Since both iterators start at the same position, call Next() once to - // advance one of the iterators. - Next(); -} - -int64_t CostAnalysisPrefetchIntervalPicker::Next() { - CHECK(!Done()) << "Prefetch interval picker's Next() is called even though " - "Done() is false"; - if (using_increasing_prefetch_time_iterator_) { - int64_t prefetch_time = increasing_prefetch_time_iterator_++; - while (increasing_prefetch_time_iterator_ <= latest_prefetch_time_ && - computation_nest_level_[increasing_prefetch_time_iterator_] != - computation_nest_level_[end_logical_time_]) { - ++increasing_prefetch_time_iterator_; - } - if (decreasing_prefetch_time_iterator_ >= earliest_prefetch_time_) { - using_increasing_prefetch_time_iterator_ = false; - } - return prefetch_time; - } else { - int64_t prefetch_time = decreasing_prefetch_time_iterator_--; - // As a compilation time optimization, reduce the number of intervals that - // this prefetch interval picker returns. When we run out of the increasing - // prefetch time iterator, only explore up to - // kNumExploredDecreasingIntervals intervals. To do that, calculate the - // 1/kNumExploredDecreasingIntervals of the elapsed time between the - // earliest prefetch time and the use, and decrement the iterator until the - // prefetch elapsed time is at least as large as this target value. This - // allows us to reduce the number of expensive heap fit and resource checks - // when the graph consists of a large number of fast-executing HLOs. - // - // Shown pictorially, assuming kNumExploredDecreasingIntervals = 3 and the - // numbers indicating the elapsed time of the HLOs, only the indicated - // options for prefetch start time would be explored: - // - // ---1---1---3---1---1---1---1---0---0---0---0---1---5---X - // ^ ^ ^ ^ - // Option3 Option2 Option1 Use - // (Earliest) - float next_target_interval_elapsed = 0; - if (increasing_prefetch_time_iterator_ > latest_prefetch_time_) { - next_target_interval_elapsed = - GetLogicalIntervalElapsed(prefetch_time, end_logical_time_) + - (GetLogicalIntervalElapsed(earliest_prefetch_time_, - end_logical_time_) / - kNumExploredDecreasingIntervals); - VLOG(3) << "Next target interval elapsed: " - << next_target_interval_elapsed; - } - while (decreasing_prefetch_time_iterator_ >= earliest_prefetch_time_ && - (computation_nest_level_[decreasing_prefetch_time_iterator_] != - computation_nest_level_[end_logical_time_] || - GetLogicalIntervalElapsed(decreasing_prefetch_time_iterator_, - end_logical_time_) < - next_target_interval_elapsed)) { - --decreasing_prefetch_time_iterator_; - } - if (increasing_prefetch_time_iterator_ <= latest_prefetch_time_) { - using_increasing_prefetch_time_iterator_ = true; - } - return prefetch_time; - } -} - -bool CostAnalysisPrefetchIntervalPicker::Done() const { - return increasing_prefetch_time_iterator_ > latest_prefetch_time_ && - decreasing_prefetch_time_iterator_ < earliest_prefetch_time_; -} - -int64_t CostAnalysisPrefetchIntervalPicker::latest_time() const { - return latest_prefetch_time_; -} - -void CostAnalysisPrefetchIntervalPicker::SetRetryNumber(int retry_number) { - retry_number_ = retry_number; -} - -int CostAnalysisPrefetchIntervalPicker::GetMinWhileNestLevel( - int64_t start_time, int64_t end_time) const { - int min_nest_level = - std::min(while_nest_level_[start_time], while_nest_level_[end_time]); - int change_idx = while_nest_level_change_[end_time]; - while (change_idx >= start_time) { - min_nest_level = std::min(min_nest_level, while_nest_level_[change_idx]); - change_idx = while_nest_level_change_[change_idx]; - } - return min_nest_level; -} - -float CostAnalysisPrefetchIntervalPicker::GetLogicalIntervalElapsed( - int64_t start_time, int64_t end_time) const { - CHECK_LE(start_time, end_time); - if (start_time == end_time) { - return 0.0; - } - if (start_time < 0) { - start_time = 0; - } - // Since elapsed_time_cumsum_ is already weighed by the while loop nesting - // level, normalize the elapsed time by dividing with the nesting factor of - // the interval (start and end times). - int interval_while_nest_level = GetMinWhileNestLevel(start_time, end_time); - return (elapsed_time_cumsum_[end_time - 1] - - elapsed_time_cumsum_[start_time]) / - while_execution_counts_[interval_while_nest_level]; -} - -std::string CostAnalysisPrefetchIntervalPicker::ToDebugString() const { - int current_logical_prefetch_time = using_increasing_prefetch_time_iterator_ - ? increasing_prefetch_time_iterator_ - : decreasing_prefetch_time_iterator_; - float logical_interval_elapsed = 0.0f; - if (current_logical_prefetch_time <= end_logical_time_) { - logical_interval_elapsed = GetLogicalIntervalElapsed( - current_logical_prefetch_time, end_logical_time_); - } - return absl::StrCat( - "Async copy elapsed (s) = ", async_copy_elapsed_, - ", inst elapsed reduction (s) = ", inst_elapsed_reduction_, - ", logical interval elapsed (s) = ", logical_interval_elapsed, - ", interval = (", current_logical_prefetch_time, ", ", end_logical_time_, - ")"); -} - -std::string CostAnalysisPrefetchIntervalPicker::ToNoCopyDebugString( - const Shape& shape, int64_t start_time, int64_t end_time) const { - float async_copy_elapsed = cost_analysis_.GetAsyncCopyElapsed( - size_override_ ? *size_override_ - : cost_analysis_.GetShapeSizeBytes(shape)); - float logical_interval_elapsed = - GetLogicalIntervalElapsed(start_time, end_time); - return absl::StrCat( - "Async copy elapsed (s) = ", async_copy_elapsed, - ", logical interval elapsed (s) = ", logical_interval_elapsed); -} - -std::optional -CostAnalysisPrefetchIntervalPicker::BufferIntervalAlternateMemoryBenefit( - const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval) - const { - return cost_analysis_.GetMemoryBoundedness(interval); -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h deleted file mode 100644 index 56eb842d1e8ebe..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h +++ /dev/null @@ -1,297 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_PREFETCH_INTERVAL_PICKER_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_PREFETCH_INTERVAL_PICKER_H_ - -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo.pb.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/shape.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -// Abstract base class that memory space assignment uses to pick prefetch -// intervals. -class PrefetchIntervalPicker { - public: - PrefetchIntervalPicker() = default; - virtual ~PrefetchIntervalPicker() = default; - - // Returns true if the buffer can be allocated in alternate memory space - // without any copies (prefetches). - virtual bool CanAllocateInAlternateMemoryNoCopy(int64_t shape_size, - int64_t start_time, - int64_t end_time) const = 0; - - // Returns the preferred end time for an eviction that starts at a given time - // and must end by the given end time. - virtual int64_t PreferredEvictionEndTime(const Shape& shape, - int64_t start_time, - int64_t latest_end_time) const = 0; - - // Returns the latest time that a prefetch can start. - virtual int64_t LatestPrefetchStartTime(int64_t shape_size, - int64_t start_time, int64_t end_time, - const HloUse* use) const = 0; - - // Returns the preferred time that a prefetch can start. - virtual int64_t PreferredPrefetchStartTime( - const Shape& shape, int64_t earliest_prefetch_start_time, - int64_t latest_prefetch_start_time, int64_t prefetch_end_time) const = 0; - - // Returns the latest time that a prefetch can end that is less than or equal - // to proposed_prefetch_end_time. - virtual int64_t LatestPrefetchEndTime( - int64_t original_prefetch_end_time, - int64_t proposed_prefetch_end_time) const { - return proposed_prefetch_end_time; - } - - // Returns the estimated end time of a prefetch that starts at the given time. - virtual int64_t EstimatedPrefetchEndTime(const Shape& shape, - int64_t start_time, - int64_t end_time) const = 0; - - // Returns the elapsed time in seconds between the logical interval that - // corresponds to the instruction schedule. - virtual float GetLogicalIntervalElapsed(int64_t start_time, - int64_t end_time) const = 0; - - // Begins the iterator for the first start time of the prefetch. - virtual void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time, - bool strict_timing = false) = 0; - - // Advances the start time of the prefetch and returns that value. - virtual int64_t Next() = 0; - - // Returns true if the available prefetch intervals have been exhausted. - virtual bool Done() const = 0; - - // Returns the latest time the prefetch interval picker will have pick. - virtual int64_t latest_time() const = 0; - - // The retry number can be used to modify the interval picking policies. The - // first attempt will have a retry_number of 0, then 1, etc. - virtual void SetRetryNumber(int retry_number) { - retry_number_ = retry_number; - } - int retry_number() const { return retry_number_; } - - // Returns a debug string for the current state of the prefetch interval - // picker. - virtual std::string ToDebugString() const = 0; - - // Returns a debug string for no-copy allocation. - virtual std::string ToNoCopyDebugString(const Shape& shape, - int64_t start_time, - int64_t end_time) const = 0; - - // Prefetch interval pickers may return a value corresponding to the benefit - // of placing the BufferInterval in the alternate memory. The larger value, - // the more beneficial. - virtual std::optional BufferIntervalAlternateMemoryBenefit( - const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval) - const { - return std::nullopt; - } - - protected: - const absl::flat_hash_map* - instruction_schedule_ = nullptr; - int retry_number_ = 0; -}; - -// Prefetch interval picker that uses instruction count to overlap asynchronous -// copies with independent computation. The min and max overlap counts describe -// the number of independent HLOs overlapped while a value is being prefetched -// into the alternate memory (between CopyStart and CopyDone HLO instructions). -// max_overlap_count attempts to prevent bringing tensors into the alternate -// memory too eagerly and hence occupying the space for other tensors which -// might use it. min_overlap_count attempts to prevent cases where tensors are -// prefetched into the alternate memory without sufficient time for the copy to -// take place. In those cases, it's just better to keep the tensor in the -// default memory instead of hurting the critical path with this copy that -// likely won't finish in time. -class InstructionCountPrefetchIntervalPicker : public PrefetchIntervalPicker { - public: - InstructionCountPrefetchIntervalPicker(int64_t min_overlap_count, - int64_t max_overlap_count) - : min_overlap_count_(min_overlap_count), - max_overlap_count_(max_overlap_count) {} - - bool CanAllocateInAlternateMemoryNoCopy(int64_t shape_size, - int64_t start_time, - int64_t end_time) const override; - - int64_t PreferredEvictionEndTime(const Shape& shape, int64_t start_time, - int64_t latest_end_time) const override; - - int64_t LatestPrefetchStartTime(int64_t shape_size, int64_t start_time, - int64_t end_time, - const HloUse* use) const override; - - int64_t PreferredPrefetchStartTime(const Shape& shape, - int64_t earliest_prefetch_start_time, - int64_t latest_prefetch_start_time, - int64_t prefetch_end_time) const override; - - int64_t EstimatedPrefetchEndTime(const Shape& shape, int64_t start_time, - int64_t end_time) const override; - float GetLogicalIntervalElapsed(int64_t start_time, - int64_t end_time) const override; - - void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time, - bool strict_timing = false) override; - - int64_t Next() override; - bool Done() const override; - - int64_t latest_time() const override; - - std::string ToDebugString() const override; - std::string ToNoCopyDebugString(const Shape& shape, int64_t start_time, - int64_t end_time) const override; - - private: - int64_t min_overlap_count_; - int64_t max_overlap_count_; - int64_t end_time_; - int64_t current_prefetch_time_; - bool strict_timing_ = false; - int64_t strict_prefetch_time_ = -1; -}; - -// Prefetch interval picker that uses cost analysis to overlap asynchronous -// copies with independent computation. It uses min (independent computation -// duration) / (asynchronous copy duration) ratio to guide whether the prefetch -// is within the lower bound. For the upper bound, it restricts the maximum -// duration that a buffer may occupy the alternate memory space as a multiple of -// the time it would take to copy a buffer that is the size of the alternate -// memory. It starts with the preferred ratio in Begin() and works its way for -// alternately earlier and later prefetches until hitting min and max ratios. -// The value for buffer size for max async copy is a mechanism to prevent -// copying small buffers between the two memories unnecessarily. For calculating -// the max time that the buffer can reside in alternate memory, we use the -// larger of this value and the actual size of the buffer. A shape override can -// also be provided which causes the interval picker to use that shape for async -// copy durations instead of the actual shape of the copy. -class CostAnalysisPrefetchIntervalPicker : public PrefetchIntervalPicker { - public: - CostAnalysisPrefetchIntervalPicker( - const CostAnalysis& cost_analysis, float min_overlap_to_async_copy_ratio, - float preferred_overlap_to_async_copy_ratio, - float max_overlap_to_mem_size_async_copy_ratio, int64_t mem_size_bytes, - std::optional size_override = std::nullopt); - - bool CanAllocateInAlternateMemoryNoCopy(int64_t shape_size, - int64_t start_time, - int64_t end_time) const override; - - int64_t PreferredEvictionEndTime(const Shape& shape, int64_t start_time, - int64_t latest_end_time) const override; - - int64_t LatestPrefetchEndTime( - int64_t original_prefetch_end_time, - int64_t proposed_prefetch_end_time) const override; - - int64_t LatestPrefetchStartTime(int64_t shape_size, int64_t start_time, - int64_t end_time, - const HloUse* use) const override; - - int64_t PreferredPrefetchStartTime(const Shape& shape, - int64_t earliest_prefetch_start_time, - int64_t latest_prefetch_start_time, - int64_t prefetch_end_time) const override; - - int64_t EstimatedPrefetchEndTime(const Shape& shape, int64_t start_time, - int64_t end_time) const override; - float GetLogicalIntervalElapsed(int64_t start_time, - int64_t end_time) const override; - - void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time, - bool strict_timing = false) override; - - int64_t Next() override; - bool Done() const override; - - int64_t latest_time() const override; - - void SetRetryNumber(int retry_number) override; - - std::string ToDebugString() const override; - std::string ToNoCopyDebugString(const Shape& shape, int64_t start_time, - int64_t end_time) const override; - - std::optional BufferIntervalAlternateMemoryBenefit( - const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval) - const override; - - private: - // Finds the minimum nest level in the given interval. - int GetMinWhileNestLevel(int64_t start_time, int64_t end_time) const; - - // Given the elapsed time to copy this buffer to the alternate memory, returns - // the longest time that this buffer may reside in the alternate memory space. - float GetMaxElapsedInAlternateMemory(float async_copy_elapsed) const; - - // For each instruction in the flattened schedule, maintain their elapsed time - // (in cumulative sum) and while nesting level. - std::vector elapsed_time_cumsum_; - std::vector while_nest_level_; - std::vector computation_nest_level_; - // Maintain the index of the most recent (before this instruction) nest level - // change in order to efficiently determine the minimum nest level in an - // interval. - std::vector while_nest_level_change_; - - const CostAnalysis& cost_analysis_; - float min_overlap_to_async_copy_ratio_; - float preferred_overlap_to_async_copy_ratio_; - float max_async_copy_elapsed_; - float async_copy_elapsed_; - float inst_elapsed_reduction_; - int64_t end_logical_time_; - int64_t earliest_prefetch_time_; - int64_t latest_prefetch_time_; - bool using_increasing_prefetch_time_iterator_ = true; - int64_t increasing_prefetch_time_iterator_; - int64_t decreasing_prefetch_time_iterator_; - - std::vector while_execution_counts_; - // Size override is used to override the size of the shape of the async copy - // to treat all async copies the same duration. Having an override forces - // prefetches to be scheduled roughly in FIFO order. - std::optional size_override_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_PREFETCH_INTERVAL_PICKER_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker_test.cc b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker_test.cc deleted file mode 100644 index d2ba1cb08bbe24..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker_test.cc +++ /dev/null @@ -1,453 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/prefetch_interval_picker.h" - -#include -#include - -#include -#include -#include "absl/log/log.h" -#include "absl/strings/string_view.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/service/memory_space_assignment/testing_utils.h" -#include "xla/shape.h" -#include "xla/shape_util.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -using CostAnalysisPrefetchIntervalPickerTest = HloHardwareIndependentTestBase; - -TEST_F(CostAnalysisPrefetchIntervalPickerTest, PrefetchIntervalOrder) { - absl::string_view hlo_string = R"( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - a = f32[2,4] negate(param0) - b = f32[2,4] negate(a) - c = f32[2,4] negate(b) - d = f32[2,4] negate(c) - e = f32[2,4] negate(d) - f = f32[2,4] negate(e) - g = f32[2,4] negate(f) - h = f32[2,4] negate(g) - i = f32[2,4] negate(h) - j = f32[2,4] negate(i) - k = f32[2,4] negate(j) - l = f32[2,4] negate(k) - m = f32[2,4] negate(l) - n = f32[2,4] negate(m) - o = f32[2,4] negate(n) - p = f32[2,4] negate(o) - q = f32[2,4] negate(p) - r = f32[2,4] negate(q) - s = f32[2,4] negate(r) - t = f32[2,4] negate(s) - u = f32[2,4] negate(t) - ROOT v = f32[2,4] add(u, param0) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - HloCostAnalysis hlo_cost_analysis; - auto hlo_cost_analysis_wrapper = - std::make_unique(hlo_cost_analysis); - auto op_cost_manager = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - CostAnalysisOptions options; - ASSERT_OK_AND_ASSIGN( - auto cost_analysis, - FakeCostAnalysis::Create(*op_cost_manager, *module, options)); - CostAnalysisPrefetchIntervalPicker interval_picker( - *cost_analysis, - /*min_overlap_to_async_copy_ratio=*/1.0, - /*preferred_overlap_to_async_copy_ratio=*/2.0, - /*max_overlap_to_mem_size_async_copy_ratio=*/4.0, - /*mem_size_bytes=*/32); - - HloInstruction* root = module->entry_computation()->root_instruction(); - const HloUse use{root, /*operand_number=*/1, /*operand_index=*/{}}; - interval_picker.Begin(use, /*start_time=*/0, /*end_time=*/22, std::nullopt); - - // Expect that the first interval is (15, 22), which has elapsed time of 6.0, - // twice of the async copy elased (3.0). Then we expect that intervals will be - // visited in alternating increasing and decreasing orders until hitting the - // min and max async copy overlap ratios, which are the intervals (18, 22) - // and (9, 22) respectively. - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 15); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 16); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 14); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 17); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 13); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 18); // Min async overlap ratio reached. - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 12); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 11); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 10); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 9); // Max async overlap ratio reached. - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_TRUE(interval_picker.Done()); - - // Expect that if the time between start_time and end_time is too short, there - // won't be any available intervals. - interval_picker.Begin(use, /*start_time=*/19, /*end_time=*/22, std::nullopt); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_TRUE(interval_picker.Done()); -} - -TEST_F(CostAnalysisPrefetchIntervalPickerTest, PrefetchIntervalOrderWhile) { - absl::string_view hlo_string = R"( - HloModule bug, is_scheduled=true - - while_condition { - param1 = (f32[2,4]) parameter(0) // 19 - ROOT cond = pred[] constant(true) // 20 - } - - while_body { - param2 = (f32[2,4]) parameter(0) // 21 - gte2 = f32[2,4] get-tuple-element(param2), index=0 // 22 - add = f32[2,4] add(gte2, gte2) // 23 - ROOT tuple2 = (f32[2,4]) tuple(add) // 24 - } - - ENTRY Entry { - param0 = f32[2,4] parameter(0) // 0 - a = f32[2,4] negate(param0) // 1 - b = f32[2,4] negate(a) // 2 - c = f32[2,4] negate(b) // 3 - d = f32[2,4] negate(c) // 4 - e = f32[2,4] negate(d) // 5 - f = f32[2,4] negate(e) // 6 - g = f32[2,4] negate(f) // 7 - h = f32[2,4] negate(g) // 8 - i = f32[2,4] negate(h) // 9 - j = f32[2,4] negate(i) // 10 - k = f32[2,4] negate(j) // 11 - l = f32[2,4] negate(k) // 12 - m = f32[2,4] negate(l) // 13 - n = f32[2,4] negate(m) // 14 - o = f32[2,4] negate(n) // 15 - p = f32[2,4] negate(o) // 16 - q = f32[2,4] negate(p) // 17 - tuple = (f32[2,4]) tuple(q) // 18 - while = (f32[2,4]) while(tuple), condition=while_condition, body=while_body // 25 - gte1 = f32[2,4] get-tuple-element(while), index=0 // 26 - r = f32[2,4] negate(gte1) // 27 - s = f32[2,4] negate(r) // 28 - t = f32[2,4] negate(s) // 29 - u = f32[2,4] negate(t) // 30 - ROOT v = f32[2,4] add(u, param0) // 31 - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - HloCostAnalysis hlo_cost_analysis; - auto hlo_cost_analysis_wrapper = - std::make_unique(hlo_cost_analysis); - auto op_cost_manager = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - CostAnalysisOptions options; - ASSERT_OK_AND_ASSIGN( - auto cost_analysis, - FakeCostAnalysis::Create(*op_cost_manager, *module, options)); - CostAnalysisPrefetchIntervalPicker interval_picker( - *cost_analysis, - /*min_overlap_to_async_copy_ratio=*/1.0, - /*preferred_overlap_to_async_copy_ratio=*/2.0, - /*max_overlap_to_mem_size_async_copy_ratio=*/12.0, - /*mem_size_bytes=*/32); - - EXPECT_EQ(cost_analysis->GetWhileNestMultiplier(1), 5.0); - HloInstruction* root = module->entry_computation()->root_instruction(); - const HloUse use{root, /*operand_number=*/1, /*operand_index=*/{}}; - interval_picker.Begin(use, /*start_time=*/0, /*end_time=*/31, std::nullopt); - - // Because there are while loop computations between [19, 24], we ensure that - // the interval picker avoids this interval. - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 25); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 26); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 18); - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 27); // Min async overlap ratio reached. - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_EQ(interval_picker.Next(), 17); // Max async overlap ratio reached. - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_TRUE(interval_picker.Done()); -} - -TEST_F(CostAnalysisPrefetchIntervalPickerTest, NestedWhile) { - // This test is to check against a bug where we didn't assign - // while_nest_level_ for while instructions, and defaulting to 0. This could - // cause the prefetch interval logic to think a nested while instruction is - // the same level as the outermost computation. - absl::string_view hlo_string = R"( - HloModule bug, is_scheduled=true - - while_condition.2 { - param1 = (f32[2,4]) parameter(0) // 11 - ROOT cond = pred[] constant(true) // 12 - } - - while_body.2 { - param2 = (f32[2,4]) parameter(0) // 13 - gte2 = f32[2,4] get-tuple-element(param2), index=0 // 14 - add = f32[2,4] add(gte2, gte2) // 15 - ROOT tuple2 = (f32[2,4]) tuple(add) // 16 - } - - while_condition.1 { - param3 = (f32[2,4]) parameter(0) // 5 - ROOT cond = pred[] constant(true) // 6 - } - - while_body.1 { - param4 = (f32[2,4]) parameter(0) // 7 - gte1 = f32[2,4] get-tuple-element(param4), index=0 // 8 - add1 = f32[2,4] add(gte1, gte1) // 9 - tuple1 = (f32[2,4]) tuple(add1) // 10 - while = (f32[2,4]) while(tuple1), condition=while_condition.2, body=while_body.2 // 17 - gte2 = f32[2,4] get-tuple-element(while), index=0 // 18 - add2 = f32[2,4] add(gte2, gte2) // 19 - ROOT tuple2 = (f32[2,4]) tuple(add2) // 20 - } - - ENTRY Entry { - param0 = f32[2,4] parameter(0) // 0 - a = f32[2,4] negate(param0) // 1 - b = f32[2,4] negate(a) // 2 - c = f32[2,4] negate(b) // 3 - tuple = (f32[2,4]) tuple(c) // 4 - while = (f32[2,4]) while(tuple), condition=while_condition.1, body=while_body.1 // 21 - gte1 = f32[2,4] get-tuple-element(while), index=0 // 22 - ROOT root = f32[2,4] add(gte1, param0) // 23 - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - CostAnalysisOptions options; - HloCostAnalysis hlo_cost_analysis; - auto hlo_cost_analysis_wrapper = - std::make_unique(hlo_cost_analysis); - auto op_cost_manager = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - ASSERT_OK_AND_ASSIGN( - auto cost_analysis, - FakeCostAnalysis::Create(*op_cost_manager, *module, options)); - CostAnalysisPrefetchIntervalPicker interval_picker( - *cost_analysis, - /*min_overlap_to_async_copy_ratio=*/1.0, - /*preferred_overlap_to_async_copy_ratio=*/2.0, - /*max_overlap_to_mem_size_async_copy_ratio=*/12.0, - /*mem_size_bytes=*/32); - - HloInstruction* root = module->entry_computation()->root_instruction(); - const HloUse use{root, /*operand_number=*/1, /*operand_index=*/{}}; - const Shape& shape = root->operand(1)->shape(); - - // We expect the root's latest prefetch start time to be before the while loop - // (logical time 4). - EXPECT_EQ(interval_picker.LatestPrefetchStartTime( - cost_analysis->GetShapeSizeBytes(shape), /*start_time=*/0, - /*end_time=*/23, &use), - 4); -} - -TEST_F(CostAnalysisPrefetchIntervalPickerTest, ConsecutiveConditionals) { - // This is a test for b/170668492, where prefetching for consecutive - // conditionals can cause the prefetch to start in the conditional's - // computation. - absl::string_view hlo_string = R"( - HloModule bug, is_scheduled=true - - true_computation.0 { - p0 = (f32[3]{0}) parameter(0) // 5 - gte = f32[3]{0} get-tuple-element(p0), index=0 // 6 - ROOT neg1 = f32[3]{0} negate(gte) // 7 - } - - false_computation.0 { - p0 = (f32[3]{0}) parameter(0) // 8 - gte = f32[3]{0} get-tuple-element(p0), index=0 // 9 - ROOT neg2 = f32[3]{0} negate(gte) // 10 - } - - true_computation.1 { - p0 = (f32[3]{0}) parameter(0) // 12 - gte = f32[3]{0} get-tuple-element(p0), index=0 // 13 - ROOT neg1 = f32[3]{0} negate(gte) // 14 - } - - false_computation.1 { - p0 = (f32[3]{0}) parameter(0) // 15 - gte = f32[3]{0} get-tuple-element(p0), index=0 // 16 - ROOT neg2 = f32[3]{0} negate(gte) // 17 - } - - ENTRY entry { - p0 = f32[3]{0} parameter(0) // 0 - p1 = f32[3]{0} parameter(1) // 1 - p2 = pred[] parameter(2) // 2 - tuple0 = (f32[3]{0}) tuple(p0) // 3 - tuple1 = (f32[3]{0}) tuple(p1) // 4 - conditional0 = f32[3]{0} conditional(p2, tuple0, tuple0), true_computation=true_computation.0, false_computation=false_computation.0 // 11 - conditional1 = f32[3]{0} conditional(p2, tuple1, tuple1), true_computation=true_computation.1, false_computation=false_computation.1 // 18 - ROOT tuple2 = (f32[3]{0}, f32[3]{0}) tuple(conditional0, conditional1) // 19 - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - CostAnalysisOptions options; - HloCostAnalysis hlo_cost_analysis; - auto hlo_cost_analysis_wrapper = - std::make_unique(hlo_cost_analysis); - auto op_cost_manager = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - ASSERT_OK_AND_ASSIGN( - auto cost_analysis, - FakeCostAnalysis::Create(*op_cost_manager, *module, options)); - CostAnalysisPrefetchIntervalPicker interval_picker( - *cost_analysis, - /*min_overlap_to_async_copy_ratio=*/1.0, - /*preferred_overlap_to_async_copy_ratio=*/2.0, - /*max_overlap_to_mem_size_async_copy_ratio=*/12.0, - /*mem_size_bytes=*/32); - - LOG(INFO) << module->ToString(); - - HloInstruction* conditional1 = - module->entry_computation()->GetInstructionWithName("conditional1"); - const HloUse use{conditional1, /*operand_number=*/1, /*operand_index=*/{0}}; - const Shape& shape = - module->entry_computation()->parameter_instruction(0)->shape(); - - // Expect that the prefetch to start before conditional0's called - // computations. - EXPECT_LT(interval_picker.LatestPrefetchStartTime( - cost_analysis->GetShapeSizeBytes(shape), /*start_time=*/0, - /*end_time=*/11, &use), - 5); -} - -TEST_F(CostAnalysisPrefetchIntervalPickerTest, EarliestLatestWindowTooSmall) { - // This tests the scenario where there is an op that takes a long time (tanh - // in this example) and as a result the earliest and latest times both fall - // inside this long-running op. In this case, we should still return a valid - // prefetch interval just before the long-running op. - absl::string_view hlo_string = R"( - HloModule bug, is_scheduled=true - - ENTRY Entry { - param0 = f32[2,4] parameter(0) - negate = f32[2,4] negate(param0) - tanh = f32[2,4] tanh(param0) - ROOT add = f32[2,4] add(tanh, negate) - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); - - CostAnalysisOptions options; - HloCostAnalysis hlo_cost_analysis; - auto hlo_cost_analysis_wrapper = - std::make_unique(hlo_cost_analysis); - auto op_cost_manager = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper), - /*enable_cache=*/false)); - ASSERT_OK_AND_ASSIGN( - auto cost_analysis, - FakeCostAnalysis::Create(*op_cost_manager, *module, options)); - cost_analysis->SetOverrideForGetInstructionElapsed( - [](const HloInstruction& hlo) { - if (hlo.opcode() == HloOpcode::kTanh) { - return 20.0; - } - return 1.0; - }); - CostAnalysisPrefetchIntervalPicker interval_picker( - *cost_analysis, - /*min_overlap_to_async_copy_ratio=*/1.0, - /*preferred_overlap_to_async_copy_ratio=*/2.0, - /*max_overlap_to_mem_size_async_copy_ratio=*/12.0, - /*mem_size_bytes=*/32); - - HloInstruction* root = module->entry_computation()->root_instruction(); - const HloUse use{root, /*operand_number=*/1, /*operand_index=*/{}}; - interval_picker.Begin(use, /*start_time=*/1, /*end_time=*/3, std::nullopt); - - LOG(INFO) << interval_picker.ToDebugString(); - EXPECT_FALSE(interval_picker.Done()); - EXPECT_EQ(interval_picker.Next(), 1); - EXPECT_TRUE(interval_picker.Done()); -} - -} // namespace -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/repacking.h b/third_party/xla/xla/service/memory_space_assignment/repacking.h deleted file mode 100644 index 0327792738ac03..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/repacking.h +++ /dev/null @@ -1,49 +0,0 @@ -/* Copyright 2020 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_REPACKING_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_REPACKING_H_ - -#include - -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/service/heap_simulator/allocation_block.h" - -namespace xla { -namespace memory_space_assignment { - -// An interface to define allocation repacking algorithms. -class MemorySpaceAssignmentRepacker { - public: - MemorySpaceAssignmentRepacker(int64_t max_size, int64_t alignment) - : max_size_(max_size), alignment_(alignment) {} - virtual ~MemorySpaceAssignmentRepacker() = default; - - // Repack the AllocationBlocks provided in the parameter. Returns true if - // allocations have been modified and false if not. Returns a non-ok status if - // there was an error. - virtual absl::StatusOr Repack( - absl::Span allocations) = 0; - - protected: - int64_t max_size_; - int64_t alignment_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_REPACKING_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/simulator.cc b/third_party/xla/xla/service/memory_space_assignment/simulator.cc deleted file mode 100644 index 0ab7fa507c3f26..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/simulator.cc +++ /dev/null @@ -1,450 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/simulator.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/strings/str_cat.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/hlo_alias_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/memory_space_assignment/allocation.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { - -void RuntimeSimulator::InitializeAlternateMemoryMap( - const AllocationSequence& allocations) { - outputs_in_alternate_memory_map_.clear(); - operands_in_alternate_memory_map_.clear(); - for (auto& allocation : allocations) { - if (allocation->is_scoped_allocation()) { - continue; - } - - if (!allocation->is_copy_allocation()) { - if (allocation->memory_space() == MemorySpace::kAlternate) { - const HloInstruction* defining_instruction = - allocation->defining_position().instruction; - outputs_in_alternate_memory_map_[defining_instruction].push_back( - allocation->defining_position().index); - } - } - for (auto& hlo_use : allocation->uses()) { - const HloInstruction* use_instruction = hlo_use.instruction; - operands_in_alternate_memory_map_[use_instruction].push_back( - std::make_pair(hlo_use.operand_number, hlo_use.operand_index)); - } - } -} - -float RuntimeSimulator::SimulateElapsedTimeWithoutAsyncCopyLikes( - const HloLiveRange& hlo_live_range, const AllocationSequence& allocations) { - InitializeAlternateMemoryMap(allocations); - const auto& instruction_sequence = - hlo_live_range.flattened_instruction_sequence().instructions(); - float total_elapsed = 0.0; - for (const HloInstruction* instruction : instruction_sequence) { - if (instruction->opcode() == HloOpcode::kWhile) { - continue; - } - - absl::Span outputs_in_alternate_memory; - auto output_it = outputs_in_alternate_memory_map_.find(instruction); - if (output_it != outputs_in_alternate_memory_map_.end()) { - outputs_in_alternate_memory = absl::MakeSpan(output_it->second); - } - - absl::Span> - operands_in_alternate_memory; - auto operand_it = operands_in_alternate_memory_map_.find(instruction); - if (operand_it != operands_in_alternate_memory_map_.end()) { - operands_in_alternate_memory = absl::MakeSpan(operand_it->second); - } - - float instruction_elapsed_per_invoke = - cost_analysis_->GetInstructionElapsedInAlternateMemory( - *instruction, operands_in_alternate_memory, - outputs_in_alternate_memory); - float total_trip_count = cost_analysis_->CalculateNestTripCount( - instruction, &cost_analysis_cache_); - // Calculate total elapsed time by summing up the overall elapsed time of - // each instruction. - total_elapsed += total_trip_count * instruction_elapsed_per_invoke; - } - return total_elapsed; -} - -bool IsAsyncCopyLikeStart(const HloInstruction* instruction) { - return instruction->opcode() == HloOpcode::kCopyStart || - (instruction->opcode() == HloOpcode::kAsyncStart && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice); -} - -bool IsAsyncCopyLikeDone(const HloInstruction* instruction) { - return (instruction->opcode() == HloOpcode::kCopyDone || - (instruction->opcode() == HloOpcode::kAsyncDone && - instruction->async_wrapped_instruction()->opcode() == - HloOpcode::kSlice)); -} - -MemoryTransferDirection GetAsyncCopyLikeDirection( - const HloInstruction* async_copy_like_start, - int64_t alternate_memory_space) { - CHECK(IsAsyncCopyLikeStart(async_copy_like_start)); - - int64_t operand_memory_space = - async_copy_like_start->operand(0)->shape().layout().memory_space(); - // Get all users - std::optional output_memory_space; - for (const HloInstruction* user : async_copy_like_start->users()) { - if (user->opcode() == HloOpcode::kCopyDone || - user->opcode() == HloOpcode::kAsyncDone) { - output_memory_space.emplace(user->shape().layout().memory_space()); - break; - } - } - if (!output_memory_space.has_value()) { - return MemoryTransferDirection::kUnsupported; - } - - if (operand_memory_space == xla::Layout::kDefaultMemorySpace && - output_memory_space == alternate_memory_space) { - return MemoryTransferDirection::kDefaultToAlternate; - } - if (operand_memory_space == alternate_memory_space && - output_memory_space == xla::Layout::kDefaultMemorySpace) { - return MemoryTransferDirection::kAlternateToDefault; - } - return MemoryTransferDirection::kUnsupported; -} - -const std::list& -RuntimeSimulator::GetOutstandingReadDefaultQueue() const { - return outstanding_read_default_queue_; -} - -const std::list& -RuntimeSimulator::GetOutstandingWriteDefaultQueue() const { - return outstanding_write_default_queue_; -} - -const HloInstruction* RuntimeSimulator::RemoveBytesFromQueueIfNotEmpty( - std::list& async_copy_like_queue, - float processed_bytes) { - if (async_copy_like_queue.empty()) return nullptr; - CHECK_GE(async_copy_like_queue.front().remaining_bytes_to_transfer, - processed_bytes); - async_copy_like_queue.front().remaining_bytes_to_transfer -= processed_bytes; - if (async_copy_like_queue.front().remaining_bytes_to_transfer == 0.0) { - const HloInstruction* retired_instruction = - async_copy_like_queue.front().copy_like_start_inst; - async_copy_like_queue.pop_front(); - return retired_instruction; - } - return nullptr; -} - -float RuntimeSimulator::SimulateAsyncCopyLikeDone( - const HloInstruction* copy_like_done_instruction) { - const HloInstruction* copy_like_start_instruction = - copy_like_done_instruction->operand(0); - MemoryTransferDirection direction = GetAsyncCopyLikeDirection( - copy_like_start_instruction, alternate_memory_space_); - if (direction == MemoryTransferDirection::kUnsupported) { - // The memory access is not a default <-> alternate memory copy. - VLOG(1) << "Unsupported memory transfer direction for copy-done: " - << copy_like_done_instruction->ToString(); - return 0.0; - } - std::list& same_direction_queue = - direction == MemoryTransferDirection::kDefaultToAlternate - ? outstanding_read_default_queue_ - : outstanding_write_default_queue_; - std::list& opposite_direction_queue = - direction == MemoryTransferDirection::kDefaultToAlternate - ? outstanding_write_default_queue_ - : outstanding_read_default_queue_; - - if (absl::c_find_if(same_direction_queue, - [&](const OutstandingAsyncCopyLike& async_copy_like) { - return async_copy_like.copy_like_start_inst == - copy_like_start_instruction; - }) == same_direction_queue.end()) { - // The copy has already finished; thus, the copy-done takes no time. - return 0.0; - } - - // Each iteration of the while loop simulates transferring a number of - // bytes from each queue that is equal to the smaller of the two elements - // at the front of each queue. If that causes us to finish a copy in the - // same_direction_queue, and that copy is the copy_like_done_instruction, we - // break the loop. - float elapsed_time = 0.0; - const HloInstruction* retired_instruction_in_same_direction_queue = nullptr; - // Loop until we process the copy start instruction that the copy-done - // instruction is waiting for. - do { - float bytes_to_process = - same_direction_queue.front().remaining_bytes_to_transfer; - float available_bandwidth = - cost_analysis_->DefaultMemBandwidthBytesPerSecond(); - - if (!opposite_direction_queue.empty()) { - // Need to share the bandwidth with the opposite direction queue. - available_bandwidth *= 0.5; - bytes_to_process = std::min( - bytes_to_process, - opposite_direction_queue.front().remaining_bytes_to_transfer); - } - - elapsed_time += bytes_to_process / available_bandwidth; - - RemoveBytesFromQueueIfNotEmpty(opposite_direction_queue, bytes_to_process); - retired_instruction_in_same_direction_queue = - RemoveBytesFromQueueIfNotEmpty(same_direction_queue, bytes_to_process); - } while (retired_instruction_in_same_direction_queue != - copy_like_start_instruction); - return elapsed_time; -}; - -RuntimeSimulator::ElapsedAndIdleTimes -RuntimeSimulator::SimulateComputeInstruction( - const HloInstruction* instruction, - absl::Span> - operands_in_alternate_memory, - absl::Span outputs_in_alternate_memory) { - // Calculate the time in which the instruction does not access the default - // memory. - float default_memory_idle_time = - cost_analysis_->GetDefaultMemoryBandwidthIdleTime( - *instruction, operands_in_alternate_memory, - outputs_in_alternate_memory); - - // Execute the outstanding async copy likes in the idle time. - default_memory_idle_time = - ProcessAsyncCopyLikesInIdleTime(default_memory_idle_time); - - float inst_elapsed = cost_analysis_->GetInstructionElapsedInAlternateMemory( - *instruction, operands_in_alternate_memory, outputs_in_alternate_memory); - return {inst_elapsed, default_memory_idle_time}; -} - -float RuntimeSimulator::ProcessAsyncCopyLikesInIdleTime(float time) { - if (time <= 0.0) { - return 0.0; - } - - double available_bandwidth = - cost_analysis_->DefaultMemBandwidthBytesPerSecond(); - - float remaining_simulation_time = time; - // This loop simulates the execution of the front memory requests in the - // read and/or write queues. The loop terminates when the remaining time is - // exhausted or there are no more outstanding async copy likes. - while ((!outstanding_read_default_queue_.empty() || - !outstanding_write_default_queue_.empty()) && - remaining_simulation_time > 0.0) { - if (!outstanding_read_default_queue_.empty() && - !outstanding_write_default_queue_.empty()) { - // Need to share the bandwidth - available_bandwidth *= 0.5; - } - float bytes_to_process = available_bandwidth * remaining_simulation_time; - if (!outstanding_read_default_queue_.empty()) { - bytes_to_process = std::min( - bytes_to_process, - outstanding_read_default_queue_.front().remaining_bytes_to_transfer); - } - if (!outstanding_write_default_queue_.empty()) { - bytes_to_process = std::min( - bytes_to_process, - outstanding_write_default_queue_.front().remaining_bytes_to_transfer); - } - - float real_elapsed_time = bytes_to_process / available_bandwidth; - remaining_simulation_time -= real_elapsed_time; - if (remaining_simulation_time <= 0.0) { - // This can happen due to floating point errors. - remaining_simulation_time = 0.0; - } - - RemoveBytesFromQueueIfNotEmpty(outstanding_read_default_queue_, - bytes_to_process); - RemoveBytesFromQueueIfNotEmpty(outstanding_write_default_queue_, - bytes_to_process); - } - - return remaining_simulation_time; -} - -namespace { - -float GetUnusedDefaultMemBandwidthBytes(float bytes_per_second, float seconds) { - CHECK_GE(bytes_per_second, 0.0); - - return bytes_per_second * seconds; -} - -} // namespace - -float RuntimeSimulator::SimulateElapsedTime( - const HloModule* hlo_module, const HloAliasAnalysis& alias_analysis, - const AllocationSequence& allocations, - const std::vector* alt_mem_bytes_occupied) { - InitializeAlternateMemoryMap(allocations); - - std::unique_ptr hlo_live_range = - HloLiveRange::Run(hlo_module->schedule(), alias_analysis, - hlo_module->entry_computation()) - .value(); - - // Cannot provide a valid result if the bandwidth is invalid. - CHECK_GT(cost_analysis_->DefaultMemBandwidthBytesPerSecond(), 0.0); - - float total_elapsed = 0.0; - // The number of additional bytes that could be transferred between default - // and alternate memory. - float cumulative_available_transfer_bytes = 0.0; - - if (alt_mem_bytes_occupied) { - CHECK_EQ( - alt_mem_bytes_occupied->size(), - hlo_live_range->flattened_instruction_sequence().instructions().size()); - } - const auto& instruction_sequence = - hlo_live_range->flattened_instruction_sequence().instructions(); - for (int time = 0; time < instruction_sequence.size(); ++time) { - const HloInstruction* instruction = instruction_sequence[time]; - float inst_elapsed = 0.0; - float idle_default_memory_bandwidth_time = 0.0; - if (instruction->opcode() == HloOpcode::kWhile) { - // Since the instructions in the while body are calculated - // separately, we can skip the while instruction. - continue; - } - if (instruction->parent()->IsAsyncComputation()) { - // We assume the overhead of async computations can be hidden perfectly. - continue; - } - if (IsAsyncCopyLikeStart(instruction)) { - // Try to categorize the async copy instruction into - // read-from-default and write-to-default queues. - MemoryTransferDirection direction = - GetAsyncCopyLikeDirection(instruction, alternate_memory_space_); - const Shape& transfer_shape = - (instruction->opcode() == HloOpcode::kCopyStart) - ? instruction->operand(0)->shape() - : ShapeUtil::GetSubshape(instruction->shape(), - /*index=*/{1}); - float transfer_bytes = - static_cast(cost_analysis_->GetShapeSizeBytes(transfer_shape)); - if (direction == MemoryTransferDirection::kDefaultToAlternate) { - outstanding_read_default_queue_.push_back( - OutstandingAsyncCopyLike{instruction, transfer_bytes}); - } else if (direction == MemoryTransferDirection::kAlternateToDefault) { - outstanding_write_default_queue_.push_back( - OutstandingAsyncCopyLike{instruction, transfer_bytes}); - } else { - // The copy does not involve default memory. - } - } else if (IsAsyncCopyLikeDone(instruction)) { - inst_elapsed = SimulateAsyncCopyLikeDone(instruction); - } else { - // This branch is for the compute instructions. - absl::Span outputs_in_alternate_memory; - auto output_it = outputs_in_alternate_memory_map_.find(instruction); - if (output_it != outputs_in_alternate_memory_map_.end()) { - outputs_in_alternate_memory = absl::MakeSpan(output_it->second); - } - - absl::Span> - operands_in_alternate_memory; - auto operand_it = operands_in_alternate_memory_map_.find(instruction); - if (operand_it != operands_in_alternate_memory_map_.end()) - operands_in_alternate_memory = absl::MakeSpan(operand_it->second); - - ElapsedAndIdleTimes elapsed_and_idle = - SimulateComputeInstruction(instruction, operands_in_alternate_memory, - outputs_in_alternate_memory); - inst_elapsed = elapsed_and_idle.elapsed_time; - idle_default_memory_bandwidth_time = - elapsed_and_idle.idle_default_memory_bandwidth_time; - } - float total_trip_count = 0.0; - if (inst_elapsed > 0.0) { - // The calculation assumes all instructions are executed independently. - // Thus, the execution time is the same for each invocation. This property - // is not hold for all cases. For example, if an async copies are - // outstanding before the loop, and there are other async copies inside - // the loop body. In this case, the first async copy in the first - // iteration will be slower than other iterations, since it needs to wait - // for the async copies issued before the loop. - total_trip_count = cost_analysis_->CalculateNestTripCount( - instruction, &cost_analysis_cache_); - total_elapsed += inst_elapsed * total_trip_count; - } - - cumulative_available_transfer_bytes += - (GetUnusedDefaultMemBandwidthBytes( - cost_analysis_->DefaultMemBandwidthBytesPerSecond(), - idle_default_memory_bandwidth_time) * - total_trip_count); - VLOG(2) << [&]() { - std::string instruction_name(instruction->name()); - if (instruction->opcode() == HloOpcode::kCopyStart && - instruction->cross_program_prefetch_index().has_value()) { - absl::StrAppend(&instruction_name, " (xprogram prefetch)"); - } - std::string alt_mem_bytes_occupied_str = ""; - if (alt_mem_bytes_occupied) { - alt_mem_bytes_occupied_str = - absl::StrCat("; alt mem usage: ", alt_mem_bytes_occupied->at(time)); - } - - return absl::StrCat(time, ": instruction: ", instruction_name, - "; elapsed: ", inst_elapsed, - "; cumulative available transfer bytes: ", - cumulative_available_transfer_bytes, - "; trip count: ", total_trip_count, - alt_mem_bytes_occupied_str); - }(); - } - - return total_elapsed; -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/simulator.h b/third_party/xla/xla/service/memory_space_assignment/simulator.h deleted file mode 100644 index 9a466733441001..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/simulator.h +++ /dev/null @@ -1,184 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SIMULATOR_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SIMULATOR_H_ - -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/shape_util.h" - -namespace xla { -namespace memory_space_assignment { - -enum class MemoryTransferDirection { - kUnsupported, - kDefaultToAlternate, - kAlternateToDefault, -}; - -// REQUIRES: -// * async_copy_like_start must be an async copy-start or slice-start -// instruction. -MemoryTransferDirection GetAsyncCopyLikeDirection( - const HloInstruction* async_copy_like_start, - int64_t alternate_memory_space); - -// This struct is used to track the outstanding async copy like instructions and -// the remaining bytes required to be accessed. -struct OutstandingAsyncCopyLike { - const HloInstruction* copy_like_start_inst; - float remaining_bytes_to_transfer; - bool operator==(const OutstandingAsyncCopyLike& other) const { - return copy_like_start_inst == other.copy_like_start_inst && - remaining_bytes_to_transfer == other.remaining_bytes_to_transfer; - } -}; - -// A wrapper class around runtime simulator. -class RuntimeSimulator { - public: - // A struct that captures an instructions elapsed time and the amount of time - // we estimate default memory bandwidth to be idle, during that instruction. - struct ElapsedAndIdleTimes { - float elapsed_time; - float idle_default_memory_bandwidth_time; - }; - - explicit RuntimeSimulator(CostAnalysis* cost_analysis, - int64_t alternate_memory_space) - : cost_analysis_(cost_analysis), - alternate_memory_space_(alternate_memory_space) {} - - // This constructor is used to inject the outstanding async copy queues for - // testing purpose. - explicit RuntimeSimulator( - CostAnalysis* cost_analysis, int64_t alternate_memory_space, - const std::list& outstanding_read_default_queue, - const std::list& - outstanding_write_default_queue) - : cost_analysis_(cost_analysis), - alternate_memory_space_(alternate_memory_space), - outstanding_read_default_queue_(outstanding_read_default_queue), - outstanding_write_default_queue_(outstanding_write_default_queue) {} - - ~RuntimeSimulator() = default; - - // This function provides a basic estimate without considering the overhead of - // async copies. - float SimulateElapsedTimeWithoutAsyncCopyLikes( - const HloLiveRange& hlo_live_range, - const AllocationSequence& allocations); - - // Returns the time to simulate the hlo_live_range, when we account for the - // waiting time for async copy like instructions to finish. - // - // To simulate the overhead of async copy like instructions, we need to - // maintain two queues to track the outstanding memory access requests that - // read/write the default memory. When we simulate compute, we use any time - // there is spare bandwidth to simulate async memory accesses to default - // memory. If we get to an async copy like done, we must wait until it - // finishes (potentially waiting for copies issued before it to finish. - // - // alt_mem_bytes_occupied is a vector of the amount of alt mem bytes allocated - // at any given instruction. It may be null. - float SimulateElapsedTime( - const HloModule* hlo_module, const HloAliasAnalysis& alias_analysis, - const AllocationSequence& allocations, - const std::vector* alt_mem_bytes_occupied = nullptr); - - // This is an auxiliary function for simulating the execution - // time for executing a copy-done instruction. It returns the - // elapsed time (in seconds) for executing the copy-done instruction. - // - // This function also updates the passed in queues as we complete async copy - // like instructions during the simulation. - // - // We simulate the shared bandwidth for default-alternate memory access. - // For example, if the copy-done instruction is a default-write memory - // process, and there are outstanding default-read memory processes in the - // outstanding_read_default_queue, then we use half of the bandwidth to - // process both requests in parallel. Otherwise, we use the full bandwidth to - // process the default-write request. - float SimulateAsyncCopyLikeDone( - const HloInstruction* copy_like_done_instruction); - - const std::list& GetOutstandingReadDefaultQueue() - const; - - const std::list& GetOutstandingWriteDefaultQueue() - const; - - // This is an auxiliary function for simulating the execution - // time for a compute instruction. It returns the elapsed time (in seconds) - // for executing the compute instruction. - // - // Aside from returning the elapsed time, this function also updates the - // outstanding memory request queues, by draining them when the compute - // instruction is not occupying bandwidth. - ElapsedAndIdleTimes SimulateComputeInstruction( - const HloInstruction* compute_instruction, - absl::Span> - operands_in_alternate_memory, - absl::Span outputs_in_alternate_memory); - - private: - // This function parses the memory space assignment solution and initializes - // the maps that record, for each instruction, which outputs and operands are - // stored in alternate memory. These maps are used to estimate the runtime of - // the HLO module. - void InitializeAlternateMemoryMap(const AllocationSequence& allocations); - const CostAnalysis* cost_analysis_; - CostAnalysis::Cache cost_analysis_cache_; - // Members used for memory model simulation - // This function updates the queue by updating the front request with the - // processed bytes. If the request is completed (no remaining bytes to - // process), the function returns the instruction and pop it from the queue. - // Otherwise, it returns nullptr. - const HloInstruction* RemoveBytesFromQueueIfNotEmpty( - std::list& async_copy_like_queue, - float processed_bytes); - - // This is an auxiliary function which simulates the process of draining - // the memory access queues in a given amount of time (seconds). If both - // outstanding_*_default_queues are non-empty, they share bandwidth. If one of - // the queues is empty and the other is not, it gets the full bandwdith. - // - // Returns the remaining idle time after processing async-copy-likes. - float ProcessAsyncCopyLikesInIdleTime(float time); - - int64_t alternate_memory_space_; - std::list outstanding_read_default_queue_; - std::list outstanding_write_default_queue_; - absl::flat_hash_map> - outputs_in_alternate_memory_map_; - absl::flat_hash_map>> - operands_in_alternate_memory_map_; -}; - -} // namespace memory_space_assignment -} // namespace xla -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SIMULATOR_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/simulator_test.cc b/third_party/xla/xla/service/memory_space_assignment/simulator_test.cc deleted file mode 100644 index c4a07ccea64af1..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/simulator_test.cc +++ /dev/null @@ -1,636 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/simulator.h" - -#include -#include -#include -#include - -#include -#include -#include "absl/container/flat_hash_map.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/status/status_macros.h" -#include "absl/strings/string_view.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/allocation.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/shape.h" -#include "tsl/platform/errors.h" - -namespace xla { -namespace { - -using memory_space_assignment::CostAnalysis; -using memory_space_assignment::CostAnalysisOptions; -using memory_space_assignment::RuntimeSimulator; - -using ::testing::ElementsAreArray; -using ::testing::IsEmpty; - -constexpr int64_t kAlternateMemorySpace = 1; - -class MemorySpaceAssignmentSimulatorTest - : public HloHardwareIndependentTestBase { - protected: - absl::Status Initialize(absl::string_view hlo_string) { - ABSL_ASSIGN_OR_RETURN(module_, ParseAndReturnVerifiedModule(hlo_string)); - for (HloInstruction* inst : module_->entry_computation()->instructions()) { - instruction_map_[inst->name()] = inst; - // Construct an allocation for the instruction if it is in the alternate - // memory. - if (inst->shape().has_layout() && - inst->shape().layout().memory_space() == kAlternateMemorySpace) { - auto allocation = - std::make_unique( - HloPosition{inst, {}}, - memory_space_assignment::MemorySpace::kAlternate, - HeapSimulator::Chunk::FromOffsetSize(-1, -1), - /*start_time=*/0, - /*end_time=*/1); - for (HloInstruction* user : inst->users()) { - allocation->AddUse(HloUse{user, 0}); - } - allocations_.push_back(std::move(allocation)); - } - } - HloCostAnalysis::Options tpu_device_options; - // Assume 1 FLOP per second for testing. - tpu_device_options.set_flops_per_second(1); - // Assume 1 byte per second for testing. - tpu_device_options.set_bytes_per_second(1); - hlo_cost_analysis_ = std::make_unique(tpu_device_options); - hlo_cost_analysis_wrapper_ = - std::make_unique(*hlo_cost_analysis_); - op_cost_manager_ = std::make_unique( - OpCostManager::Options{ - /*enable_cache=*/false, - /*enable_analysis_logging=*/false, - }, - OpCostManager::CalculationNode::CreateLeaf( - "HloCostAnalysis", - CreateHloCostAnalysisCalculator(*hlo_cost_analysis_wrapper_), - /*enable_cache=*/false)); - CostAnalysisOptions cost_analysis_options; - // Assume 2 byte per second for testing. - cost_analysis_options.alternate_mem_read_bandwidth_bytes_per_second = 2; - cost_analysis_options.alternate_mem_write_bandwidth_bytes_per_second = 2; - cost_analysis_options.default_mem_bandwidth_bytes_per_second = 1.0; - - ABSL_ASSIGN_OR_RETURN(alias_analysis_, - HloAliasAnalysis::Run(module_.get(), &alias_info_)); - ABSL_ASSIGN_OR_RETURN( - cost_analysis_, - CostAnalysis::Create(*op_cost_manager_, cost_analysis_options, - &alias_info_, *module_, alias_analysis_.get())); - ABSL_ASSIGN_OR_RETURN(hlo_live_range_, - HloLiveRange::Run(module_->schedule(), *alias_analysis_, - module_->entry_computation())); - runtime_simulator_ = std::make_unique( - cost_analysis_.get(), kAlternateMemorySpace); - return absl::OkStatus(); - } - absl::flat_hash_map - instruction_map_; - std::unique_ptr hlo_cost_analysis_; - std::unique_ptr hlo_cost_analysis_wrapper_; - std::unique_ptr op_cost_manager_; - std::unique_ptr cost_analysis_; - std::unique_ptr alias_analysis_; - std::unique_ptr hlo_live_range_; - memory_space_assignment::AllocationSequence allocations_; - std::unique_ptr runtime_simulator_; - std::unique_ptr module_; - AliasInfo alias_info_; -}; - -TEST_F(MemorySpaceAssignmentSimulatorTest, SingleLayerLoop) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - - %body { - %constant.1 = s32[] constant(1) - %param = (s32[]) parameter(0) - %count = s32[] get-tuple-element(%param), index=0 - %increment = s32[] add(s32[] %count, s32[] %constant.1) - ROOT %loop_result = (s32[]) tuple(%increment) - } - - %condition { - %param = (s32[]) parameter(0) - %constant.42 = s32[] constant(42) - %condition_input = s32[] get-tuple-element(%param), index=0 - ROOT %greater = pred[] compare(s32[] %constant.42, s32[] %condition_input), direction=GT - } - - ENTRY Entry { - %dummy_input = s32[] parameter(0) - %constant.0 = s32[] constant(0) - ROOT %while = (s32[]) while(tuple(%constant.0)), condition=%condition, body=%body - } - - )"; - ASSERT_OK(Initialize(hlo_string)); - // The total elapsed time is the summation of the elapsed time of each - // instruction. Here are the overhead of each instruction (secs): - // %increment: 12 * 42 - // tuple(%constant.0): 8 * 1 - // %greater: 9 * 42 - // %loop_result: 8 * 42 - EXPECT_EQ(runtime_simulator_->SimulateElapsedTimeWithoutAsyncCopyLikes( - *hlo_live_range_, allocations_), - 1226); - EXPECT_EQ(runtime_simulator_->SimulateElapsedTime( - module_.get(), *alias_analysis_, allocations_), - 1226); -} - -TEST_F(MemorySpaceAssignmentSimulatorTest, NestedLayerLoop) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - %inner.body { - %constant.1 = s32[] constant(1) - %param = (s32[]) parameter(0) - %count = s32[] get-tuple-element(%param), index=0 - %increment = s32[] add(s32[] %count, s32[] %constant.1) - ROOT %loop_result = (s32[]) tuple(%increment) - } - %inner.condition { - %param = (s32[]) parameter(0) - %constant.42 = s32[] constant(42) - %condition_input = s32[] get-tuple-element(%param), index=0 - ROOT %greater = pred[] compare(s32[] %constant.42, s32[] %condition_input), direction=GT - } - %outer.body { - %constant.0 = s32[] constant(0) - %constant.1 = s32[] constant(1) - %param = (s32[]) parameter(0) - %inner_while = (s32[]) while(tuple(%constant.0)), condition=%inner.condition, body=%inner.body - %count = s32[] get-tuple-element(%param), index=0 - %increment = s32[] add(s32[] %count, s32[] %constant.1) - ROOT %loop_result = (s32[]) tuple(%increment) - } - %outer.condition { - %param = (s32[]) parameter(0) - %constant.27 = s32[] constant(27) - %condition_input = s32[] get-tuple-element(%param), index=0 - ROOT %greater = pred[] compare(s32[] %constant.27, s32[] %condition_input), direction=GT - } - ENTRY Entry { - %constant.0 = s32[] constant(0) - ROOT %while_outer = (s32[]) while(tuple(%constant.0)), condition=%outer.condition, body=%outer.body - } - )"; - ASSERT_OK(Initialize(hlo_string)); - // The inner loop is derived from the SingleLayerLoop test, whose overhead is - // 1226 seconds. - - // For the outer loop, the overhead of each instruction is: - // %increment: 12 * 27 - // tuple(%constant.0): 8 * 1 - // %greater: 9 * 27 - // %loop_result: 8 * 27 - // Thus, the total overhead of the while_outer is 1226 * 27 + 12 * 27 + 8 * 1 - // + 9 * 27 + 8 * 27 = 33893 - - EXPECT_EQ(runtime_simulator_->SimulateElapsedTimeWithoutAsyncCopyLikes( - *hlo_live_range_, allocations_), - 33893); - EXPECT_EQ(runtime_simulator_->SimulateElapsedTime( - module_.get(), *alias_analysis_, allocations_), - 33893); -} - -TEST_F(MemorySpaceAssignmentSimulatorTest, SingleAsyncCopyOverhead) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[1,1,1024,2048] parameter(0) - copy-start.1 = (f32[1,1,1024,2048]{0,1,2,3:S(1)}, f32[1,1,1024,2048], u32[]) copy-start(param_0) - ROOT copy-done.1 = f32[1,1,1024,2048]{0,1,2,3:S(1)} copy-done(copy-start.1) - } - - )"; - ASSERT_OK(Initialize(hlo_string)); - - // Since the HLO does not contain memory access, pass an empty allocation - // sequence for test. - memory_space_assignment::AllocationSequence allocations; - // The SimulateElapsedTimeWithoutAsyncCopyLikes should not include the - // overhead of async copies. - EXPECT_EQ(runtime_simulator_->SimulateElapsedTimeWithoutAsyncCopyLikes( - *hlo_live_range_, allocations_), - 0); - // The expected elapsed time is 1024 * 2048 * 4 / 1 = 8388608. - EXPECT_EQ(runtime_simulator_->SimulateElapsedTime( - module_.get(), *alias_analysis_, allocations_), - 8388608); -} - -TEST_F(MemorySpaceAssignmentSimulatorTest, AsyncCopyWithComputationOverhead) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[8] parameter(0) - param_1 = f32[2] parameter(1) - copy-start.1 = (f32[8]{0:S(1)}, f32[8], u32[]) copy-start(param_0) - neg_compute = f32[2] negate(param_1) - ROOT copy-done.1 = f32[8]{0:S(1)} copy-done(copy-start.1) - } - - )"; - ASSERT_OK(Initialize(hlo_string)); - // The neg_compute read/write 16 bytes in total, thus, it requires 16 seconds - // for default memory access. Since it only requires 2 FLOPs computation which - // requires 2 seconds, it is a memory-bound instruction which does not have - // idle time to process async copies. - // Workflow: - // neg_compute: | 16 sec (memory-bound) | - // copy-done.1: | | read 32 bytes | - // time: | 16 sec | 32 sec | - EXPECT_EQ(runtime_simulator_->SimulateElapsedTime( - module_.get(), *alias_analysis_, allocations_), - 48); -} - -TEST_F(MemorySpaceAssignmentSimulatorTest, SingleAsyncSliceCopyOverhead) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[3072,2048] parameter(0) - slice-start = ((f32[3072,2048]), f32[768,2048]{1,0:S(1)}, s32[]) slice-start(f32[3072,2048] param_0), slice={[1536:2304], [0:2048]} - ROOT slice-done = f32[768,2048]{1,0:T(8,128)S(1)} slice-done(((f32[3072,2048]), f32[768,2048]{1,0:S(1)}, s32[]) slice-start) - } - )"; - ASSERT_OK(Initialize(hlo_string)); - - memory_space_assignment::AllocationSequence allocations; - // The expected elapsed time is 768 * 2048 * 4 / 1 = 6291456. - float expected_elapsed_time = 6291456; - - EXPECT_EQ(runtime_simulator_->SimulateElapsedTime( - module_.get(), *alias_analysis_, allocations_), - expected_elapsed_time); -} - -TEST_F(MemorySpaceAssignmentSimulatorTest, - AsyncCopyAndAsyncSliceAndComputeOverhead) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[2048] parameter(0) - param_1 = f32[64] parameter(1) - param_2 = f32[128] parameter(2) - slice-start = ((f32[2048]), f32[64]{0:S(1)}, s32[]) slice-start(f32[2048] param_0), slice={[0:64]} - copy-start = (f32[64]{0:S(1)}, f32[64], u32[]) copy-start(f32[64] param_1) - slice-done = f32[64]{0:S(1)} slice-done(((f32[2048]), f32[64]{0:S(1)}, s32[]) slice-start) - copy-done = f32[64]{0:S(1)} copy-done(copy-start) - copy-start-overlap = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(f32[128] param_2) - add = f32[64]{0:S(1)} add(slice-done, copy-done) - ROOT copy-done-overlap = f32[128]{0:S(1)} copy-done(copy-start-overlap) - } - )"; - ASSERT_OK(Initialize(hlo_string)); - - // The overhead of each instruction is: - // slice-done: 64 * 4 / 1 = 256 sec (default memory access) - // copy-done: 64 * 4 /1 = 256 sec (default memory access) - // add: 3 * 64 * 4 / 2 = 384 sec (alternate memory access) - // Since add does not access default memory, we can use process 384 bytes in - // copy-start-overlap. - // copy-done-overlap: (128 * 4 - 384) / 1 = 128 sec (default memory access) - EXPECT_EQ(runtime_simulator_->SimulateElapsedTime( - module_.get(), *alias_analysis_, allocations_), - 1024); -} - -class SimulateAsyncCopyLikeDoneTest - : public MemorySpaceAssignmentSimulatorTest { - protected: - absl::Status Initialize(absl::string_view hlo_string) { - ABSL_RETURN_IF_ERROR(MemorySpaceAssignmentSimulatorTest::Initialize(hlo_string)); - if (instruction_map_.contains("copy-start.1")) { - outstanding_read_default_queue_.push_back( - memory_space_assignment::OutstandingAsyncCopyLike{ - instruction_map_["copy-start.1"], 512}); - } - if (instruction_map_.contains("copy-start.2")) { - outstanding_write_default_queue_.push_back( - memory_space_assignment::OutstandingAsyncCopyLike{ - instruction_map_["copy-start.2"], 128}); - } - runtime_simulator_ = std::make_unique( - cost_analysis_.get(), kAlternateMemorySpace, - outstanding_read_default_queue_, outstanding_write_default_queue_); - return absl::OkStatus(); - } - std::list - outstanding_read_default_queue_; - std::list - outstanding_write_default_queue_; -}; - -TEST_F(SimulateAsyncCopyLikeDoneTest, AsyncCopyAlreadyCompleted) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - - const HloInstruction* copy_done_inst = instruction_map_["copy-done.1"]; - // Process the copy-start.1 - runtime_simulator_->SimulateAsyncCopyLikeDone(copy_done_inst); - - // There should be no request in the read/write queues. - EXPECT_THAT(runtime_simulator_->GetOutstandingReadDefaultQueue(), IsEmpty()); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); - // The function should return 0 for requests that are already completed. - float elapsed_time_for_completed_copy = - runtime_simulator_->SimulateAsyncCopyLikeDone(copy_done_inst); - EXPECT_EQ(elapsed_time_for_completed_copy, 0); - // There should be no request in the read/write queues. - EXPECT_THAT(runtime_simulator_->GetOutstandingReadDefaultQueue(), IsEmpty()); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, AsyncCopyFullBandwidth) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - const HloInstruction* copy_done_inst = instruction_map_["copy-done.1"]; - - // The elapsed time for copy-done.1 is 128 * 4 / 1 = 512. - float copy_done_elapsed_time = - runtime_simulator_->SimulateAsyncCopyLikeDone(copy_done_inst); - EXPECT_EQ(copy_done_elapsed_time, 512); - - // There should be no request in the read/write queues. - EXPECT_THAT(runtime_simulator_->GetOutstandingReadDefaultQueue(), IsEmpty()); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, AsyncCopySharedBandwidth) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - param_1 = f32[32]{0:S(1)} parameter(1) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - copy-start.2 = (f32[32], f32[32]{0:S(1)}, u32[]) copy-start(param_1) - copy-done.2 = f32[32] copy-done(copy-start.2) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - - const HloInstruction* copy_start_1_inst = instruction_map_["copy-start.1"]; - const HloInstruction* copy_done_2_inst = instruction_map_["copy-done.2"]; - - // The copy-start.2 needs to share bandwidth with copy-start.1. Thus, it can - // only use half bandwidth to access default memory. Thus, the elapsed time is - // 32 * 4 / 0.5 = 256 - float copy_done_2_elapsed_time = - runtime_simulator_->SimulateAsyncCopyLikeDone(copy_done_2_inst); - EXPECT_EQ(copy_done_2_elapsed_time, 256); - - // The only write request (copy-start.2) should be completed. - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); - - // The read request has (128-32)*4 bytes left to process. - EXPECT_THAT( - runtime_simulator_->GetOutstandingReadDefaultQueue(), - ElementsAreArray({memory_space_assignment::OutstandingAsyncCopyLike{ - copy_start_1_inst, 384}})); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, AsyncCopyTransferPartialProcess) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - param_1 = f32[32]{0:S(1)} parameter(1) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - copy-start.2 = (f32[32], f32[32]{0:S(1)}, u32[]) copy-start(param_1) - copy-done.2 = f32[32] copy-done(copy-start.2) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - - const HloInstruction* copy_start_1_inst = instruction_map_["copy-start.1"]; - const HloInstruction* copy_done_1_inst = instruction_map_["copy-done.1"]; - const HloInstruction* copy_done_2_inst = instruction_map_["copy-done.2"]; - - // Execute copy-done.2. - float copy_done_2_elapsed_time = - runtime_simulator_->SimulateAsyncCopyLikeDone(copy_done_2_inst); - // For copy-done.2, it requires to transfer 32*4 bytes - // default-write request. At the same time, there is a 128*4 bytes - // default-read request in the queue for copy-start.1. So the - // elapsed time for copy-done.2 is 32*4 / (0.5*1) = 256. - EXPECT_EQ(copy_done_2_elapsed_time, 256); - // In parallel with copy-done.2, copy-start.1 is also being processed. - // So the remaining bytes should be 128*4 - 32*4 = 384. - EXPECT_THAT( - runtime_simulator_->GetOutstandingReadDefaultQueue(), - ElementsAreArray({memory_space_assignment::OutstandingAsyncCopyLike{ - copy_start_1_inst, 384}})); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); - - // Execute copy-done.1. - float copy_done_1_elapsed_time = - runtime_simulator_->SimulateAsyncCopyLikeDone(copy_done_1_inst); - // The copy-done.1 is the only request in the read-queue, and there is no - // request in the write-queue. Thus, it can use the full bandwidth. The - // elapsed time is 384 / 1 = 384. - EXPECT_EQ(copy_done_1_elapsed_time, 384); - // No request should be in the queue. - EXPECT_THAT(runtime_simulator_->GetOutstandingReadDefaultQueue(), IsEmpty()); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, - SimulateComputeInstructionWithSingleAsyncCopy) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - param_1 = f32[32] parameter(1) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - neg = f32[32] negate(param_1) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - const HloInstruction* copy_start_1_inst = instruction_map_["copy-start.1"]; - const HloInstruction* neg_inst = instruction_map_["neg"]; - - float compute_elapsed_time = - runtime_simulator_ - ->SimulateComputeInstruction(neg_inst, - /*operands_in_alternate_memory=*/{}, - /*outputs_in_alternate_memory=*/{}) - .elapsed_time; - - // The compute operand requires 32 FLOPs and 32 * 4 * 2 bytes access, which - // requires 32 and 256 secs respectively. Thus, it is default memory access - // dominated, which does not have idle time to process the async copy. - EXPECT_EQ(compute_elapsed_time, 256); - EXPECT_THAT( - runtime_simulator_->GetOutstandingReadDefaultQueue(), - ElementsAreArray({memory_space_assignment::OutstandingAsyncCopyLike{ - copy_start_1_inst, 512}})); - - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, - SimulateComputeInstructionWithSharedBandwidth) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - param_1 = f32[32]{0:S(1)} parameter(1) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - copy-start.2 = (f32[32], f32[32]{0:S(1)}, u32[]) copy-start(param_1) - neg = f32[32] negate(param_1) - copy-done.2 = f32[32] copy-done(copy-start.2) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - - const HloInstruction* copy_start_1_inst = instruction_map_["copy-start.1"]; - const HloInstruction* copy_start_2_inst = instruction_map_["copy-start.2"]; - - // The instruction reads 32 * 4 bytes from alternate memory, which takes 64 - // secs. In this 64 secs, it does not access default memory. Thus, we can - // process the async copies in this time. Both queues are not empty, so the - // bandwidth is shared. Each of the request at the front of the queue process - // 64 sec * 0.5 bytes/sec = 32 bytes. - float compute_elapsed_time = - runtime_simulator_ - ->SimulateComputeInstruction( - instruction_map_["neg"], - /*operands_in_alternate_memory=*/{{0, {}}}, - /*outputs_in_alternate_memory=*/{}) - .elapsed_time; - // 64 secs for alternate memory access + 128 secs for default memory access - EXPECT_EQ(compute_elapsed_time, 192); - - EXPECT_THAT( - runtime_simulator_->GetOutstandingReadDefaultQueue(), - ElementsAreArray({memory_space_assignment::OutstandingAsyncCopyLike{ - copy_start_1_inst, 480}})); - - EXPECT_THAT( - runtime_simulator_->GetOutstandingWriteDefaultQueue(), - ElementsAreArray({memory_space_assignment::OutstandingAsyncCopyLike{ - copy_start_2_inst, 96}})); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, - SimulateComputeInstructionWithFullBandwidth) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - param_1 = f32[32]{0:S(1)} parameter(1) - copy-start.1 = (f32[128]{0:S(1)}, f32[128], u32[]) copy-start(param_0) - neg = f32[32] negate(param_1) - ROOT copy-done.1 = f32[128]{0:S(1)} copy-done(copy-start.1) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - - const HloInstruction* copy_start_1_inst = instruction_map_["copy-start.1"]; - - // Same as the 'SimulateComputeInstructionWithSharedBandwidth' test, there are - // 64 secs idle time to process async copies. Since only the read queue is not - // empty, we can use the full bandwidth and process 64 sec * 1 bytes/sec = 64 - // bytes. - float compute_elapsed_time = - runtime_simulator_ - ->SimulateComputeInstruction( - instruction_map_["neg"], - /*operands_in_alternate_memory=*/{{0, {}}}, - /*outputs_in_alternate_memory=*/{}) - .elapsed_time; - // 64 secs for alternate memory access + 128 secs for default memory access - EXPECT_EQ(compute_elapsed_time, 192); - - EXPECT_THAT( - runtime_simulator_->GetOutstandingReadDefaultQueue(), - ElementsAreArray({memory_space_assignment::OutstandingAsyncCopyLike{ - copy_start_1_inst, 448}})); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); -} - -TEST_F(SimulateAsyncCopyLikeDoneTest, - SimulateComputeInstructionWithEmptyQueues) { - absl::string_view hlo_string = - R"(HloModule module, is_scheduled=true - ENTRY Entry { - param_0 = f32[128] parameter(0) - ROOT neg = f32[128] negate(param_0) - } - )"; - - ASSERT_OK(Initialize(hlo_string)); - - float compute_elapsed_time = - runtime_simulator_ - ->SimulateComputeInstruction(instruction_map_["neg"], - /*operands_in_alternate_memory=*/{}, - /*outputs_in_alternate_memory=*/{}) - .elapsed_time; - // Execution time: 128 * 4 * 2 / 1 for default access - EXPECT_EQ(compute_elapsed_time, 1024); - // The queues should remain empty. - EXPECT_THAT(runtime_simulator_->GetOutstandingReadDefaultQueue(), IsEmpty()); - EXPECT_THAT(runtime_simulator_->GetOutstandingWriteDefaultQueue(), IsEmpty()); -} - -} // namespace -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/slice.cc b/third_party/xla/xla/service/memory_space_assignment/slice.cc deleted file mode 100644 index 2869e686d8eca9..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/slice.cc +++ /dev/null @@ -1,160 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/slice.h" - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/functional/any_invocable.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/str_join.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/service/time_utils.h" -#include "xla/shape.h" - -namespace xla::memory_space_assignment { - -std::tuple -SliceDecisionToTuple(const SliceDecision& decision) { - return std::make_tuple( - std::ref(decision.chunk), decision.exclusive_start_time, - std::ref(decision.sizing), decision.copy_resource_consumed); -} - -std::string SliceDecision::ToString() const { - return absl::StrCat("{ chunk: ", chunk.ToString(), - ", (exclusive) start_time: ", exclusive_start_time, - ", sizing: ", sizing.ToString(), - ", copy_resource_consumed: ", copy_resource_consumed, - " }"); -} - -bool SliceDecision::operator==(const SliceDecision& other) const { - return SliceDecisionToTuple(*this) == SliceDecisionToTuple(other); -} - -std::string SliceProposal::ToString() const { - return absl::StrCat( - "{ slice_shape: ", slice_shape.ToString(true), ", slice_params: { ", - absl::StrJoin(slice_params, ", ", - [](std::string* out, const SliceParam& param) { - absl::StrAppend(out, param.ToString()); - }), - " }, slice_size: ", slice_size, " }"); -} - -std::ostream& operator<<(std::ostream& os, const SliceProposal& proposal) { - os << proposal.ToString(); - return os; -} - -std::tuple&, int64_t> -SliceProposal::ToTuple() const { - return std::make_tuple(std::ref(slice_shape), std::ref(slice_params), - slice_size); -} - -bool SliceProposal::operator==(const SliceProposal& other) const { - return ToTuple() == other.ToTuple(); -} - -std::string SliceParam::ToString() const { - return absl::StrCat("[", start_inclusive, ",", end_exclusive, ")"); -} - -bool SliceParam::operator==(const SliceParam& other) const { - return start_inclusive == other.start_inclusive && - end_exclusive == other.end_exclusive; -} - -bool IsUniformSliceSizingEnabled(const SlicedPrefetchOptions& options) { - return options.max_slices() > 0 && options.preferred_slice_size() > 0; -} - -std::vector SlicedPrefetchStartTimePicker::Pick( - int64_t num_slices, int64_t exclusive_prefetch_start_time, - int64_t prefetch_end_time, absl::AnyInvocable elapsed_fn, - absl::AnyInvocable has_same_parent_fn) { - CHECK_LE(exclusive_prefetch_start_time, prefetch_end_time); - VLOG(5) << "Picking slice start times. num_slices = " << num_slices - << "; exclusive_prefetch_start_time = " - << exclusive_prefetch_start_time - << "; prefetch_end_time = " << prefetch_end_time; - - // Prefetching starts after the selected start instruction and ends - // before the selected end instruction. Thus, we have (end - (start + 1)) HLO - // instructions worth of time to perform all of the sliced copies. So, the - // only choices for start times that give us time to copy are <= - // prefetch_end_time - 2. - if (exclusive_prefetch_start_time >= prefetch_end_time - 2 || - num_slices == 1) { - return std::vector(num_slices, exclusive_prefetch_start_time); - } - - float total_elapsed = - elapsed_fn(exclusive_prefetch_start_time, prefetch_end_time); - if (total_elapsed <= 0.0) { - return std::vector(num_slices, exclusive_prefetch_start_time); - } - - std::vector start_times; - start_times.reserve(num_slices); - start_times.push_back(exclusive_prefetch_start_time); - int64_t last_valid_candidate = exclusive_prefetch_start_time; - int64_t candidate = exclusive_prefetch_start_time; - while (candidate < prefetch_end_time - 1 && start_times.size() < num_slices) { - float target_elapsed = total_elapsed * - static_cast(num_slices - start_times.size()) / - static_cast(num_slices); - float elapsed = elapsed_fn(candidate, prefetch_end_time); - if (elapsed < target_elapsed) { - // We've gone past our target, so use the last valid candidate. - start_times.push_back(last_valid_candidate); - continue; - } - bool updating_candidate_impacts_elapsed = - last_valid_candidate != candidate && - elapsed_fn(last_valid_candidate, - ExclusiveToInclusiveStartTime(candidate)) > 0.0; - // has_same_parent_fn will look up the computation parent of the - // instructions at prefetch_start_time and prefetch_end_time. If - // prefetch_start_time is -1, no such instruction will exist. However, if we - // want to insert an instruction after the -1 schedule position, we can - // use the parent of the instruction at index 0 instead. Thus, we use - // std::max below. - if (has_same_parent_fn(std::max(0, exclusive_prefetch_start_time), - std::max(0, candidate)) && - updating_candidate_impacts_elapsed) { - last_valid_candidate = candidate; - } - ++candidate; - } - while (start_times.size() < num_slices) { - start_times.push_back(last_valid_candidate); - } - - return start_times; -} - -} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/memory_space_assignment/slice.h b/third_party/xla/xla/service/memory_space_assignment/slice.h deleted file mode 100644 index f0caa04e92ee41..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/slice.h +++ /dev/null @@ -1,150 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -// This file contains definitions for MSA slicing. Slicing is an allocation -// technique in which we allocate a buffer in slices that can start at different -// times, but once allocated, form a contiguous buffer. When copying buffers, we -// may want to allocate a buffer in slices, so that we delay allocating memory -// that would otherwise not be in use, due to copy bandwidth constraints. -// -// The following illustrates a buffer that is fully allocated at time t2, via -// slices. -// -// space -// ^ -// p3 | +-----------+ -// | | s2 | -// p2 | +---+-----------+ -// | | s1 | -// p1 | +-------+-------+ -// | | s0 | -// p0 | +-------+ -// +---|---|---|---|---|----> time -// t0 t1 t2 t3 t4 - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SLICE_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SLICE_H_ - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/functional/any_invocable.h" -#include "absl/status/statusor.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/shape.h" - -namespace xla::memory_space_assignment { - -// The target of a custom call that slicing uses to concatenate slices -// that are already contiguous in memory, into a larger buffer. -inline constexpr char kConcatBitcastCustomCall[] = "ConcatBitcast"; - -// The parameters for slicing a single dimension of a tensor. -struct SliceParam { - std::string ToString() const; - bool operator==(const SliceParam& other) const; - - int64_t start_inclusive; - int64_t end_exclusive; -}; - -// A proposed way to slice a buffer. -struct SliceProposal { - std::string ToString() const; - friend std::ostream& operator<<(std::ostream& os, - const SliceProposal& proposal); - std::tuple&, int64_t> ToTuple() - const; - bool operator==(const SliceProposal& other) const; - - // Shape resulting from the slice. - Shape slice_shape; - - // slice_params map to the parameters that would be passed to a slice - // instruction. Thus: - // * There should be a slice parameter for every dimension in the shape of - // the tensor being sliced. - // * The ith slice_param applies to the ith logical dimension in the shape - // being sliced. - // * If a dimension is not being sliced, it should have a SliceParam of - // {0, dim size}. - std::vector slice_params; - - // The size to be allocated for the slice. Note, this may be > the size of - // the slice shape, due to additional padding that may occur when the slices - // are concatenated back together. - int64_t slice_size; -}; - -// A SliceProposalCollection proposes a way to to slice an AllocationRequest. -// A SliceProposalCollection is generated from a SliceProposalFunction and is -// used when we want to slice a prefetch. -using SliceProposalCollection = std::vector; -using SliceProposalFunction = - std::function( - const Shape& shape, const SlicedPrefetchOptions& options)>; - -// A SliceDecision is a SliceProposal that we've determined where and when to -// allocate. -struct SliceDecision { - std::string ToString() const; - bool operator==(const SliceDecision& other) const; - - HeapSimulator::Chunk chunk; - int64_t exclusive_start_time; - SliceProposal sizing; - float copy_resource_consumed; -}; - -// Returns true if the options indicates that there is a preferred slice -// size. -bool IsUniformSliceSizingEnabled(const SlicedPrefetchOptions& options); - -// A class for turning a copy start time and end time into slice start times. -class SlicedPrefetchStartTimePicker { - public: - // Returns the amount of time elapsed in the instruction schedule between - // (exclusive_start_time, exclusive_end_time). - using ElapsedTimeFn = std::add_pointer::type; - - // Returns true if the instructions at lhs_time and rhs_time are in the same - // computation. - using SameComputationParentFn = - std::add_pointer::type; - - // Picks slice start times, given the num_slices, prefetch_start_time, and - // prefetch_end_time. The returned times are exclusive. - // - // REQUIRES: - // - The instructions following each start time are guaranateed to be in the - // same computation. - // - The returned times sorted. - // - The first returned time is equal to prefetch_start_time. - static std::vector Pick( - int64_t num_slices, int64_t exclusive_prefetch_start_time, - int64_t prefetch_end_time, absl::AnyInvocable elapsed_fn, - absl::AnyInvocable has_same_parent_fn); -}; - -} // namespace xla::memory_space_assignment - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_SLICE_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/slice_test.cc b/third_party/xla/xla/service/memory_space_assignment/slice_test.cc deleted file mode 100644 index 53fb497373e76d..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/slice_test.cc +++ /dev/null @@ -1,224 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/slice.h" - -#include -#include -#include -#include - -#include -#include -#include "absl/functional/any_invocable.h" -#include "absl/log/check.h" -#include "xla/service/time_utils.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -class SlicedPrefetchStartTimePickerTest : public ::testing::Test { - protected: - struct FakeInstructionData { - float elapsed_time = 0.0; - std::string computation; - }; - - std::vector Pick( - const std::vector& schedule_data, int64_t num_slices, - int64_t prefetch_start_time, int64_t prefetch_end_time) { - return SlicedPrefetchStartTimePicker::Pick( - num_slices, prefetch_start_time, prefetch_end_time, - [&schedule_data](int64_t exclusive_start_time, - int64_t exclusive_end_time) { - auto start_it = schedule_data.begin() + - ExclusiveToInclusiveStartTime(exclusive_start_time); - auto end_it = (exclusive_end_time < schedule_data.size() - ? schedule_data.begin() + exclusive_end_time - : schedule_data.end()); - return std::accumulate( - start_it, end_it, 0.0, - [](float total, const FakeInstructionData& data) { - return total + data.elapsed_time; - }); - }, - [&schedule_data](int64_t lhs_time, int64_t rhs_time) { - CHECK_GE(lhs_time, 0); - CHECK_GE(rhs_time, 0); - CHECK_LT(lhs_time, schedule_data.size()); - CHECK_LT(rhs_time, schedule_data.size()); - return schedule_data[lhs_time].computation == - schedule_data[rhs_time].computation; - }); - } -}; - -TEST_F(SlicedPrefetchStartTimePickerTest, Base1) { - // The 2nd slice naturally should start after 1.5 time units have passed, - // forcing us to start before t=1. - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {1.0, "a"}, - /*t=2*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/3), - ::testing::ElementsAre(-1, 0)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, Base2) { - // The 2nd slice naturally should start after 6.0 time units have passed, - // forcing us to start before t=0. - EXPECT_THAT(Pick({ - /*t=0*/ {10.0, "a"}, - /*t=1*/ {1.0, "a"}, - /*t=2*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/3), - ::testing::ElementsAre(-1, -1)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, Base3) { - // The 2nd slice naturally should start after 1.0 time unit has passed. - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/2), - ::testing::ElementsAre(-1, 0)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, Zeros1) { - // The 2nd slice naturally should start after 1.0 time unit has passed. - // Make sure we don't add extra 0.0 cost instructions to the start time. - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {0.0, "a"}, - /*t=2*/ {0.0, "a"}, - /*t=3*/ {0.0, "a"}, - /*t=4*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/5), - ::testing::ElementsAre(-1, 0)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, Zeros2) { - // The 2nd slice naturally should start after 2.0 time units have passed. - // Make sure we don't add extra 0.0 cost instructions to the start time. - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {0.0, "a"}, - /*t=2*/ {1.0, "a"}, - /*t=3*/ {0.0, "a"}, - /*t=4*/ {1.0, "a"}, - /*t=5*/ {0.0, "a"}, - /*t=6*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/7), - ::testing::ElementsAre(-1, 2)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, Zeros3) { - // The first slice always comes at prefetch_start_time. The 2nd slice - // naturally should start after 1.5 time units have passed, causing us to - // start after t=2. Make sure we don't add extra 0.0 cost instructions to the - // start time. - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {0.0, "a"}, - /*t=2*/ {1.0, "a"}, - /*t=3*/ {0.0, "a"}, - /*t=4*/ {1.0, "a"}, - /*t=5*/ {0.0, "a"}, - /*t=6*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/1, - /*prefetch_end_time=*/7), - ::testing::ElementsAre(1, 2)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, MidSchedule) { - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {1.0, "a"}, - /*t=3*/ {1.0, "a"}, - /*t=4*/ {1.0, "a"}, - /*t=5*/ {1.0, "a"}, - /*t=6*/ {1.0, "a"}, - /*t=7*/ {1.0, "a"}, - /*t=8*/ {1.0, "a"}, - /*t=9*/ {1.0, "a"}, - /*t=10*/ {1.0, "a"}, - /*t=11*/ {1.0, "a"}, - /*t=12*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/5, - /*prefetch_end_time=*/10), - ::testing::ElementsAre(5, 7)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, ManySlices) { - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {1.0, "a"}, - /*t=2*/ {1.0, "a"}, - /*t=3*/ {1.0, "a"}, - /*t=4*/ {1.0, "a"}, - /*t=5*/ {1.0, "a"}, - /*t=6*/ {1.0, "a"}, - /*t=7*/ {1.0, "a"}, - /*t=8*/ {1.0, "a"}, - /*t=9*/ {1.0, "a"}, - /*t=10*/ {1.0, "a"}, - /*t=11*/ {1.0, "a"}, - /*t=12*/ {1.0, "a"}, - /*t=13*/ {1.0, "a"}, - /*t=14*/ {1.0, "a"}, - /*t=15*/ {1.0, "a"}, - /*t=16*/ {1.0, "a"}, - /*t=17*/ {1.0, "a"}, - /*t=18*/ {1.0, "a"}, - /*t=19*/ {1.0, "a"}, - }, - /*num_slices=*/5, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/20), - ::testing::ElementsAre(-1, 3, 7, 11, 15)); -} - -TEST_F(SlicedPrefetchStartTimePickerTest, DifferentParents) { - // The 2nd slice naturally should start after t=2, but we are forced to push - // it after t=1, since the instruction at t=3 has parent "b", while the first - // instruction has parent "a." - EXPECT_THAT(Pick({ - /*t=0*/ {1.0, "a"}, - /*t=1*/ {1.0, "a"}, - /*t=2*/ {1.0, "b"}, - /*t=3*/ {1.0, "b"}, - /*t=4*/ {1.0, "b"}, - /*t=5*/ {1.0, "a"}, - }, - /*num_slices=*/2, /*prefetch_start_time=*/-1, - /*prefetch_end_time=*/6), - ::testing::ElementsAre(-1, 1)); -} - -} // namespace -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/testing_utils.h b/third_party/xla/xla/service/memory_space_assignment/testing_utils.h deleted file mode 100644 index fec767bfb815c8..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/testing_utils.h +++ /dev/null @@ -1,139 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_TESTING_UTILS_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_TESTING_UTILS_H_ - -#include -#include -#include -#include - -#include "absl/memory/memory.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_alias_analysis.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/call_graph.h" -#include "xla/service/cost_modelling/op_cost.h" -#include "xla/service/hlo_cost_analysis.h" -#include "xla/service/memory_space_assignment/cost_analysis.h" -#include "xla/shape.h" -#include "xla/shape_util.h" -#include "tsl/platform/statusor.h" - -namespace xla { -namespace memory_space_assignment { - -// For testing purposes, we define a cost analysis where we can control the -// elapsed times of each HLO and asynchronous copy. -class FakeCostAnalysis : public CostAnalysis { - public: - static absl::StatusOr> Create( - OpCostManager& op_cost_manager, const HloModule& module, - const CostAnalysisOptions& options) { - auto alias_info = std::make_unique(); - ABSL_ASSIGN_OR_RETURN(auto alias_analysis, - HloAliasAnalysis::Run(&module, alias_info.get())); - ABSL_ASSIGN_OR_RETURN(auto hlo_live_range, - HloLiveRange::Run(module.schedule(), *alias_analysis, - module.entry_computation())); - auto call_graph = CallGraph::Build(&module); - return absl::WrapUnique( - new FakeCostAnalysis(op_cost_manager, options, - std::move(alias_analysis), std::move(alias_info), - std::move(hlo_live_range), std::move(call_graph))); - } - - float GetInstructionElapsed( - const HloInstruction& instruction) const override { - if (get_instruction_elapsed_override_) { - return get_instruction_elapsed_override_(instruction); - } - return 1.0; - } - - float GetInstructionElapsedInAlternateMemory( - const HloInstruction& instruction, - absl::Span> - operands_in_alternate_mem, - absl::Span outputs_in_alternate_mem) const override { - if (get_instruction_elapsed_in_alternate_memory_override_) { - return get_instruction_elapsed_in_alternate_memory_override_( - instruction, operands_in_alternate_mem, outputs_in_alternate_mem); - } - if (!operands_in_alternate_mem.empty()) { - return 0.5; - } else { - return 1.0; - } - } - - float GetAsyncCopyElapsed(int64_t size_in_bytes) const override { - if (get_async_copy_elapsed_override_) { - return get_async_copy_elapsed_override_(size_in_bytes); - } - return 3.0; - } - - // The following methods can be used to override what the above API calls - // return. - void SetOverrideForGetInstructionElapsed( - std::function function) { - get_instruction_elapsed_override_ = function; - } - void SetOverrideForGetInstructionElapsedInAlternateMemory( - std::function>, - absl::Span)> - function) { - get_instruction_elapsed_in_alternate_memory_override_ = function; - } - void SetOverrideForGetAsyncCopyElapsed( - std::function function) { - get_async_copy_elapsed_override_ = function; - } - - protected: - FakeCostAnalysis(OpCostManager& op_cost_manager, - const CostAnalysisOptions& options, - std::unique_ptr alias_analysis, - std::unique_ptr alias_info, - std::unique_ptr hlo_live_range, - std::unique_ptr call_graph) - : CostAnalysis(op_cost_manager, options, alias_analysis.get(), - std::move(hlo_live_range), std::move(call_graph)), - alias_analysis_(std::move(alias_analysis)), - alias_info_(std::move(alias_info)) {} - - private: - std::function - get_instruction_elapsed_override_ = nullptr; - std::function>, - absl::Span)> - get_instruction_elapsed_in_alternate_memory_override_ = nullptr; - std::function get_async_copy_elapsed_override_ = nullptr; - std::unique_ptr alias_analysis_; - std::unique_ptr alias_info_; -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_TESTING_UTILS_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/tuning_utils.cc b/third_party/xla/xla/service/memory_space_assignment/tuning_utils.cc deleted file mode 100644 index cdad1b87be6e96..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/tuning_utils.cc +++ /dev/null @@ -1,64 +0,0 @@ -/* Copyright 2021 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/tuning_utils.h" - -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/log/check.h" -#include "xla/hlo/ir/hlo_module.h" - -namespace xla { - -namespace memory_space_assignment { - -void CustomizeSortedBufferInterval( - std::optional> memory_space_assignment_config, - std::vector& sorted_buffer_intervals) { - // A copy of the sorted buffer intervals to assist the creating of the - // customized buffer intervals vector respecting the config. - std::vector sorted_buffer_intervals_copy( - sorted_buffer_intervals); - - std::vector config; - if (!memory_space_assignment_config.has_value()) { - config.resize(sorted_buffer_intervals_copy.size()); - absl::c_iota(config, 0); - } else { - config = *memory_space_assignment_config; - } - - CHECK_EQ(config.size(), sorted_buffer_intervals_copy.size()); - sorted_buffer_intervals.clear(); - for (int i = 0; i < config.size(); ++i) { - sorted_buffer_intervals.push_back(sorted_buffer_intervals_copy[config[i]]); - } - - if (!sorted_buffer_intervals.empty()) { - HloModule* module = - sorted_buffer_intervals[0].buffer->instruction()->GetModule(); - - // Update the memory space assignment auto-tuning config of a module with a - // given config. - module->mutable_config().mutable_memory_space_assignment_config()->assign( - config.begin(), config.end()); - } -} - -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/tuning_utils.h b/third_party/xla/xla/service/memory_space_assignment/tuning_utils.h deleted file mode 100644 index adeee2823f5734..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/tuning_utils.h +++ /dev/null @@ -1,43 +0,0 @@ -/* Copyright 2021 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_TUNING_UTILS_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_TUNING_UTILS_H_ - -#include -#include -#include - -#include "xla/hlo/ir/hlo_module.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" - -namespace xla { - -namespace memory_space_assignment { - -using BufferInterval = - GlobalDecreasingSizeBestFitHeap::BufferInterval; - -// This function converts sorted_buffer_intervals to order-customized buffer -// intervals respecting a given memory space assignment autotuning config. -void CustomizeSortedBufferInterval( - std::optional> memory_space_assignment_config, - std::vector& sorted_buffer_intervals); - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_TUNING_UTILS_H_ diff --git a/third_party/xla/xla/service/memory_space_assignment/utils.cc b/third_party/xla/xla/service/memory_space_assignment/utils.cc deleted file mode 100644 index 363c3ac4f03364..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/utils.cc +++ /dev/null @@ -1,607 +0,0 @@ -/* Copyright 2020 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/memory_space_assignment/utils.h" - -#include -#include -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/status/status_macros.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/string_view.h" -#include "highwayhash/arch_specific.h" -#include "highwayhash/hh_types.h" -#include "highwayhash/highwayhash.h" -#include "re2/re2.h" -#include "xla/hlo/ir/hlo_casting_utils.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_instructions.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/ir/hlo_print_options.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" -#include "xla/shape_util.h" -#include "xla/util.h" - -namespace xla { -namespace memory_space_assignment { -namespace { - -// The default seed used by HloRandomFilter when no seed is given. -const int64_t kRandomFilterDefaultSeed = 1234; - -} // namespace - -bool MemorySpaceAssignmentUtils::IsInstructionOnConfiguredExecThread( - const HloInstruction& instruction, - const absl::flat_hash_set& execution_threads) { - return instruction.parent() == nullptr || execution_threads.empty() || - execution_threads.contains(instruction.parent()->execution_thread()); -} - -bool MemorySpaceAssignmentUtils::IsValueAllowedInAlternateMemory( - const HloValue* value, int64_t alternate_memory_space) { - // If the buffer is a tuple, don't use this algorithm for now. The buffers - // that are pointed to by the tuple will still use this algorithm. Because - // tuples are cheap to place in the alternate memory (they are just pointers) - // we don't need to use prefetch/evict logic. - if (value->shape().IsTuple()) { - VLOG(4) << "Keeping value " << value->ToShortString() - << " in default mem because it is a tuple."; - return false; - } - - // Don't place scalars in the alternate memory. - if (ShapeUtil::IsEffectiveScalar(value->shape())) { - VLOG(4) << "Keeping value " << value->ToShortString() - << " in default mem because it is a scalar."; - return false; - } - - // TODO(berkin): Not allocating add-dependencies either since they need to be - // treated specially. We should revisit this later. - for (const HloPosition& position : value->positions()) { - if (position.instruction->opcode() == HloOpcode::kAddDependency) { - VLOG(4) << "Keeping value " << value->ToShortString() - << " in default mem because it has a " - << "add-dependency position."; - return false; - } - } - - // Send and Recv HLOs return a request identifier. These should not be - // allocated in the alternate memory. - for (const HloPosition& position : value->positions()) { - if ((position.instruction->opcode() == HloOpcode::kSend || - position.instruction->opcode() == HloOpcode::kRecv) && - DynCast(position.instruction) - ->is_host_transfer()) { - // TODO(berkin): Host transfers using alternate memory space doesn't seem - // to work at the moment. - VLOG(4) << "Keeping value " << value->ToShortString() - << " in default mem because it is a send/recv buffer used for " - "host transfer."; - return false; - } - - // If the tensor is pre-colored to a memory space that is neither the - // default (0) nor the alternate, disallow it from the alternate memory - // space. - int64_t memory_space = 0; - if (position.shape().has_layout()) { - memory_space = position.shape().layout().memory_space(); - } - if (memory_space != 0 && memory_space != alternate_memory_space) { - VLOG(4) << "Value " << value->ToShortString() - << " not allowed in the alternate memory space due to existing " - "memory space: " - << memory_space; - return false; - } - } - - return true; -} - -bool MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory( - const GlobalDecreasingSizeBestFitHeap::BufferInterval& interval, - int64_t alternate_memory_space) { - return IsValueAllowedInAlternateMemory(interval.buffer, - alternate_memory_space) && - absl::c_all_of(interval.colocations, - [alternate_memory_space](const HloValue* value) { - return IsValueAllowedInAlternateMemory( - value, alternate_memory_space); - }); -} - -bool MemorySpaceAssignmentUtils::DoesUseMatchOperandFilter( - const HloOperandFilter& filter, const HloUse& hlo_use, - int64_t operand_size) { - // The order of checks is such that the most expensive checks are done last. - if (filter.has_size_gte() && operand_size < filter.size_gte()) { - return false; - } - if (filter.has_size_lte() && operand_size > filter.size_lte()) { - return false; - } - if (filter.has_operand_number() && - hlo_use.operand_number != filter.operand_number()) { - return false; - } - if (filter.has_tuple_index() && - hlo_use.operand_index != ShapeIndex(filter.tuple_index().index().begin(), - filter.tuple_index().index().end())) { - return false; - } - if (filter.has_instruction_name_regex() && - !RE2::FullMatch(hlo_use.instruction->name(), - filter.instruction_name_regex())) { - return false; - } - if (filter.has_instruction_regex() && - !RE2::FullMatch(hlo_use.instruction->ToString(), - filter.instruction_regex())) { - return false; - } - return true; -} - -double GetInstructionUniformRandom(const HloInstruction& instruction, - int64_t seed) { - // We use the instruction Fingerprint because it's robust across different - // runs and compilations, as well as because it depends only on the - // isomorphic computation graph of the instruction rather than the variable - // names. We would like to make identical decisions for identical - // instructions. - // - // See HloPrintOptions::Fingerprint() for details. - // - // If in the future that turns out to not be good enough, we can also pass an - // option in the proto to choose different string representations of the - // instructions - for example, the name or "instruction.ToString()" if we want - // a finer grained perturbation, or alternatively the - // "instruction.SignatureString()" if we want a coarser perturbation. - const std::string& instruction_identifier = - instruction.ToString(HloPrintOptions::Fingerprint()); - std::string instruction_seed_str = absl::StrCat( - instruction_identifier.size(), ":", instruction_identifier, ":", seed); - - // We use highwayhash as our hashing function, because it is "Strong", - // relatively fast, and "all (64/128/256 bit) variants of HighwayHash frozen, - // i.e. unchanging forever" which is good for test stability and repeatability - // of experiments. - - // PI in hex - static constexpr highwayhash::HHKey hash_key = { - 0x3243F6A8885A308Dull, - 0x313198A2E0370734ull, - 0x4A4093822299F31Dull, - 0x0082EFA98EC4E6C8ull, - }; - - highwayhash::HHStateT state(hash_key); - highwayhash::HHResult64 seed_result; - highwayhash::HighwayHashT(&state, instruction_seed_str.data(), - instruction_seed_str.size(), &seed_result); - - // Assuming highwayhash is a "good" hash and uniformly distributed across all - // 64 bits, we could probably skip the step of going through the random - // library to generate our uniform [0...1) random. - // - // We add this anyway to make it clear we're looking for a uniform random - // number. - std::mt19937_64 gen(seed_result); - return std::uniform_real_distribution()(gen); -} - -bool MemorySpaceAssignmentUtils::DoesInstructionMatchRandomFilter( - const HloPositionMatcher& filter, const HloInstruction& instruction) { - if (!filter.has_hlo_random_filter()) { - // This HloPositionMatcher doesn't have a random filter. We return that we - // aren't blocking this instruction. - return true; - } - const auto& hlo_random_filter = filter.hlo_random_filter(); - double selection_range_begin = 0.; - if (hlo_random_filter.has_selection_range_begin()) { - selection_range_begin = hlo_random_filter.selection_range_begin(); - } - double selection_range_end = 1.; - if (hlo_random_filter.has_selection_range_end()) { - selection_range_end = hlo_random_filter.selection_range_end(); - } - int64_t seed = kRandomFilterDefaultSeed; - if (hlo_random_filter.has_seed()) { - seed = hlo_random_filter.seed(); - } - double rnd = GetInstructionUniformRandom(instruction, seed); - return rnd >= selection_range_begin && rnd < selection_range_end; -} - -bool MemorySpaceAssignmentUtils::DoesPositionMatchFilter( - const HloPositionMatcher& filter, - const MsaBufferInterval& buffer_interval) { - // The order of checks is such that the most expensive checks are done last. - HloInstruction* instruction = buffer_interval.buffer->instruction(); - if (filter.has_size_gte() && filter.size_gte() > buffer_interval.size) { - return false; - } - if (filter.has_size_lte() && filter.size_lte() < buffer_interval.size) { - return false; - } - if (filter.has_tuple_index() && - buffer_interval.buffer->index() != - ShapeIndex(filter.tuple_index().index().begin(), - filter.tuple_index().index().end())) { - return false; - } - return DoesInstructionMatchFilter(filter, *instruction) && - DoesBufferIntervalMatchHloUseFilter(filter, buffer_interval) && - DoesInstructionMatchRandomFilter(filter, *instruction); -} - -bool MemorySpaceAssignmentUtils::DoesPositionMatchPositionFilter( - const HloPositionMatcher& filter, const HloPosition& position, - int64_t size) { - if (filter.has_size_gte() && filter.size_gte() > size) { - return false; - } - if (filter.has_size_lte() && filter.size_lte() < size) { - return false; - } - if (filter.has_tuple_index() && - position.index != ShapeIndex(filter.tuple_index().index().begin(), - filter.tuple_index().index().end())) { - return false; - } - return DoesInstructionMatchFilter(filter, *position.instruction) && - DoesInstructionMatchRandomFilter(filter, *position.instruction); -} - -bool MemorySpaceAssignmentUtils::DoesPositionMatchOperandFilter( - const HloOperandFilter& filter, const HloPosition& position, int64_t size) { - if (filter.has_size_gte() && filter.size_gte() > size) { - return false; - } - if (filter.has_size_lte() && filter.size_lte() < size) { - return false; - } - if (filter.has_tuple_index() && - position.index != ShapeIndex(filter.tuple_index().index().begin(), - filter.tuple_index().index().end())) { - return false; - } - if (filter.has_instruction_name_regex() && - !RE2::FullMatch(position.instruction->name(), - filter.instruction_name_regex())) { - return false; - } - if (filter.has_instruction_regex() && - !RE2::FullMatch(position.instruction->ToString(), - filter.instruction_regex())) { - return false; - } - return true; -} -bool MemorySpaceAssignmentUtils::DoesInstructionMatchFilter( - const HloPositionMatcher& filter, const HloInstruction& instruction) { - if (filter.has_instruction_name_regex() && - !RE2::FullMatch(instruction.name(), filter.instruction_name_regex())) { - return false; - } - if (filter.has_instruction_regex() && - !RE2::FullMatch(instruction.ToString(), filter.instruction_regex())) { - return false; - } - return true; -} - -bool MemorySpaceAssignmentUtils::DoesBufferIntervalMatchHloUseFilter( - const HloPositionMatcher& filter, - const MsaBufferInterval& buffer_interval) { - if (!filter.has_hlo_use_filter()) { - return true; - } - for (const HloUse& use : buffer_interval.buffer->GetUses()) { - if (DoesUseMatchOperandFilter(filter.hlo_use_filter(), use, - buffer_interval.size)) { - return true; - } - } - return false; -} - -absl::StatusOr -MemorySpaceAssignmentUtils::GetScheduleTimeFromInstructionMatcher( - const HloPositionMatcher& position_matcher, - const absl::flat_hash_map& schedule) { - for (auto schedule_entry : schedule) { - if (DoesInstructionMatchFilter(position_matcher, *schedule_entry.first)) { - return schedule_entry.second; - } - } - return NotFound("Reference instruction %s was not found in the schedule.", - position_matcher.DebugString()); -} - -absl::StatusOr> -MemorySpaceAssignmentUtils::GetPrefetchTimeByEagerness( - float prefetch_eagerness, int64_t earliest_prefetch_time, - int64_t latest_prefetch_time) { - CHECK_GE(prefetch_eagerness, 0.0); - CHECK_LE(prefetch_eagerness, 1.0); - if (earliest_prefetch_time > latest_prefetch_time) { - return static_cast>(std::nullopt); - } - return static_cast>( - earliest_prefetch_time + - (latest_prefetch_time - earliest_prefetch_time) * prefetch_eagerness); -} - -absl::StatusOr> -MemorySpaceAssignmentUtils::GetPrefetchTimeAfterInstruction( - const HloPositionMatcher& after_instruction, - const absl::flat_hash_map& schedule) { - ABSL_ASSIGN_OR_RETURN( - auto reference_instruction_time, - GetScheduleTimeFromInstructionMatcher(after_instruction, schedule)); - return static_cast>(reference_instruction_time); -} - -absl::StatusOr> -MemorySpaceAssignmentUtils::GetPrefetchTimeBeforeInstruction( - const HloPositionMatcher& before_instruction, - const absl::flat_hash_map& schedule) { - ABSL_ASSIGN_OR_RETURN( - auto reference_instruction_time, - GetScheduleTimeFromInstructionMatcher(before_instruction, schedule)); - return static_cast>(reference_instruction_time - 1); -} -absl::StatusOr> -MemorySpaceAssignmentUtils::GetPrefetchTime( - const PreferredPrefetchOverrideOptions& override_options, - int64_t earliest_prefetch_time, int64_t latest_prefetch_time, - const absl::flat_hash_map& - instruction_schedule) { - switch (override_options.options_case()) { - case PreferredPrefetchOverrideOptions::kPrefetchEagerness: - return GetPrefetchTimeByEagerness(override_options.prefetch_eagerness(), - earliest_prefetch_time, - latest_prefetch_time); - case PreferredPrefetchOverrideOptions::kAfterInstruction: - return GetPrefetchTimeAfterInstruction( - override_options.after_instruction(), instruction_schedule); - case PreferredPrefetchOverrideOptions::kBeforeInstruction: - return GetPrefetchTimeBeforeInstruction( - override_options.before_instruction(), instruction_schedule); - case PreferredPrefetchOverrideOptions::kLogicalTime: - return static_cast>( - override_options.logical_time()); - case PreferredPrefetchOverrideOptions::OPTIONS_NOT_SET: - break; - } - return static_cast>>(std::nullopt); -} - -absl::StatusOr> -MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( - const PreferredPrefetchOverrides& preferred_prefetch_overrides, - int64_t operand_size, const HloUse& hlo_use, - const absl::flat_hash_map& - instruction_schedule, - int64_t earliest_prefetch_time, int64_t latest_prefetch_time) { - for (const auto& override : preferred_prefetch_overrides.overrides()) { - if (!MemorySpaceAssignmentUtils::DoesUseMatchOperandFilter( - override.hlo_operand_filter(), hlo_use, operand_size)) { - continue; - } - VLOG(3) << "Config match for instruction " << hlo_use.instruction->name() - << " operand number " << hlo_use.operand_number << " operand index " - << hlo_use.operand_index.ToString() << " size " << operand_size - << " live range (" << earliest_prefetch_time << ", " - << latest_prefetch_time << ")"; - ABSL_ASSIGN_OR_RETURN( - auto prefetch_time, - GetPrefetchTime(override.override_options(), earliest_prefetch_time, - latest_prefetch_time, instruction_schedule)); - if (prefetch_time.has_value() && - prefetch_time.value() >= earliest_prefetch_time && - prefetch_time.value() <= latest_prefetch_time) { - return prefetch_time; - } - } - return static_cast>>(std::nullopt); -} - -absl::StatusOr> -MemorySpaceAssignmentUtils::GetPrefetchOverrideInfo( - const MsaTensorOverrides& msa_tensor_overrides, - const PreferredPrefetchOverrides& preferred_prefetch_overrides, - int64_t operand_size, const HloUse& hlo_use, - const absl::flat_hash_map& - instruction_schedule, - int64_t earliest_prefetch_time, int64_t latest_prefetch_time, - const std::optional& position) { - // First check msa_tensor_overrides. - for (const auto& override : msa_tensor_overrides.overrides()) { - if (!override.has_prefetch()) { - continue; - } - bool matches = false; - if (override.has_hlo_operand_filter() && - DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, - operand_size)) { - matches = true; - } else if (override.has_hlo_position_matcher() && position.has_value() && - DoesPositionMatchPositionFilter(override.hlo_position_matcher(), - *position, operand_size)) { - matches = true; - } - if (!matches) { - continue; - } - const auto& prefetch_options = override.prefetch(); - ABSL_ASSIGN_OR_RETURN( - auto prefetch_time, - GetPrefetchTime(prefetch_options, earliest_prefetch_time, - latest_prefetch_time, instruction_schedule)); - bool strict_timing = prefetch_options.has_strict_timing() && - prefetch_options.strict_timing(); - PrefetchOverrideInfo info; - info.prefetch_time = prefetch_time; - info.strict_timing = strict_timing; - info.fail_on_unsatisfied_override = true; - return info; - } - - // Next check preferred_prefetch_overrides. - for (const auto& override : preferred_prefetch_overrides.overrides()) { - if (!DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, - operand_size)) { - continue; - } - ABSL_ASSIGN_OR_RETURN( - auto prefetch_time, - GetPrefetchTime(override.override_options(), earliest_prefetch_time, - latest_prefetch_time, instruction_schedule)); - bool strict_timing = override.override_options().has_strict_timing() && - override.override_options().strict_timing(); - bool fail_on_unsatisfied = - override.fail_on_unsatisfied_override() || - preferred_prefetch_overrides.fail_on_unsatisfied_override() || - strict_timing; - PrefetchOverrideInfo info; - info.prefetch_time = prefetch_time; - info.strict_timing = strict_timing; - info.fail_on_unsatisfied_override = fail_on_unsatisfied; - return info; - } - return std::nullopt; -} - -bool MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( - const MsaTensorOverrides& msa_tensor_overrides, const HloPosition& position, - int64_t size) { - for (const auto& override : msa_tensor_overrides.overrides()) { - if (!override.has_pin_in_alternate_memory()) { - continue; - } - if (override.has_hlo_position_matcher() && - DoesPositionMatchPositionFilter(override.hlo_position_matcher(), - position, size)) { - return true; - } - if (override.has_hlo_operand_filter() && - DoesPositionMatchOperandFilter(override.hlo_operand_filter(), position, - size)) { - return true; - } - } - return false; -} - -bool MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - const MsaTensorOverrides& msa_tensor_overrides, const HloPosition& position, - int64_t size) { - for (const auto& override : msa_tensor_overrides.overrides()) { - if (!override.has_keep_in_default_memory()) { - continue; - } - if (override.has_hlo_position_matcher() && - DoesPositionMatchPositionFilter(override.hlo_position_matcher(), - position, size)) { - return true; - } - if (override.has_hlo_operand_filter() && - DoesPositionMatchOperandFilter(override.hlo_operand_filter(), position, - size)) { - return true; - } - } - return false; -} - -bool MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( - const MsaTensorOverrides& msa_tensor_overrides, const HloUse& hlo_use, - int64_t operand_size) { - for (const auto& override : msa_tensor_overrides.overrides()) { - if (!override.has_keep_in_default_memory()) { - continue; - } - if (override.has_hlo_operand_filter() && - DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, - operand_size)) { - return true; - } - } - return false; -} - -bool MemorySpaceAssignmentUtils::DoesCrossProgramPrefetchBufferMatchAnyFilter( - const MsaSortOrderOverrides& sort_order_overrides, - const MsaBufferInterval& buffer_interval) { - return GetBufferIntervalOverridePriority( - sort_order_overrides, buffer_interval, - /*is_cross_program_prefetch=*/true) < 0; -} - -int64_t MemorySpaceAssignmentUtils::GetBufferIntervalOverridePriority( - const MsaSortOrderOverrides& msa_sort_order_overrides, - const MsaBufferInterval& buffer_interval, bool is_cross_program_prefetch) { - if (msa_sort_order_overrides.overrides_size() == 0) { - return 0; - } - for (int64_t i = 0; i < msa_sort_order_overrides.overrides_size(); ++i) { - const auto& override = msa_sort_order_overrides.overrides(i); - if (is_cross_program_prefetch && - (!override.has_apply_to_cross_program_prefetches() || - !override.apply_to_cross_program_prefetches())) { - continue; - } - if (!MemorySpaceAssignmentUtils::DoesPositionMatchFilter( - override.hlo_position_matcher(), buffer_interval)) { - continue; - } - VLOG(3) << "Override Sort Order Config " << i << " matches " - << buffer_interval.buffer->instruction()->ToString(); - switch (override.override_options().options_case()) { - case MsaSortOrderOverrideOptions::kAssignFirst: - return std::numeric_limits::lowest() + i; - case MsaSortOrderOverrideOptions::kAssignLast: - return std::numeric_limits::max() - i; - case MsaSortOrderOverrideOptions::kAssignValue: - return override.override_options().assign_value(); - case MsaSortOrderOverrideOptions::OPTIONS_NOT_SET: - continue; - } - } - return 0; -} -} // namespace memory_space_assignment -} // namespace xla diff --git a/third_party/xla/xla/service/memory_space_assignment/utils.h b/third_party/xla/xla/service/memory_space_assignment/utils.h deleted file mode 100644 index b8a897252f3e35..00000000000000 --- a/third_party/xla/xla/service/memory_space_assignment/utils.h +++ /dev/null @@ -1,165 +0,0 @@ -/* Copyright 2020 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_UTILS_H_ -#define XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_UTILS_H_ - -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/utils/hlo_live_range.h" -#include "xla/service/heap_simulator/heap_simulator.h" -#include "xla/service/hlo_value.h" -#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" - -namespace xla { -namespace memory_space_assignment { - -using MsaBufferInterval = - GlobalDecreasingSizeBestFitHeap::BufferInterval; - -// Info resolved from prefetch overrides. -struct PrefetchOverrideInfo { - std::optional prefetch_time; - bool strict_timing = false; - bool fail_on_unsatisfied_override = false; -}; - -// Encapsulates common utility methods for memory space assignment. -class MemorySpaceAssignmentUtils { - public: - // Returns true if this buffer is allowed to be placed in the alternate - // memory. - static bool IsIntervalAllowedInAlternateMemory( - const MsaBufferInterval& interval, int64_t alternate_memory_space); - - // Returns true if the HloValue is allowed to be placed in alternate memory. - static bool IsValueAllowedInAlternateMemory(const HloValue* value, - int64_t alternate_memory_space); - - // Returns true if the instruction runs on one of the configured execution - // threads. Returns true if the set of execution threads is empty or the - // instruction has no parent computation. - static bool IsInstructionOnConfiguredExecThread( - const HloInstruction& instruction, - const absl::flat_hash_set& execution_threads); - - static bool DoesUseMatchOperandFilter(const HloOperandFilter& filter, - const HloUse& hlo_use, - int64_t operand_size); - - static bool DoesInstructionMatchFilter(const HloPositionMatcher& filter, - const HloInstruction& instruction); - - static bool DoesInstructionMatchRandomFilter( - const HloPositionMatcher& filter, const HloInstruction& instruction); - - static bool DoesPositionMatchFilter(const HloPositionMatcher& filter, - const MsaBufferInterval& buffer_interval); - - static bool DoesPositionMatchPositionFilter(const HloPositionMatcher& filter, - const HloPosition& position, - int64_t size); - - static bool DoesPositionMatchOperandFilter(const HloOperandFilter& filter, - const HloPosition& position, - int64_t size); - - static absl::StatusOr - GetScheduleTimeFromInstructionMatcher( - const HloPositionMatcher& position_matcher, - const absl::flat_hash_map& schedule); - - static absl::StatusOr> GetPrefetchTimeByEagerness( - float prefetch_eagerness, int64_t earliest_prefetch_time, - int64_t latest_prefetch_time); - - static absl::StatusOr> GetPrefetchTimeAfterInstruction( - const HloPositionMatcher& after_instruction, - const absl::flat_hash_map& schedule); - - static absl::StatusOr> - GetPrefetchTimeBeforeInstruction( - const HloPositionMatcher& before_instruction, - const absl::flat_hash_map& schedule); - - static absl::StatusOr> GetPrefetchTime( - const PreferredPrefetchOverrideOptions& override_options, - int64_t earliest_prefetch_time, int64_t latest_prefetch_time, - const absl::flat_hash_map& - instruction_schedule); - - static absl::StatusOr> - GetOverriddenPreferredPrefetchTime( - const PreferredPrefetchOverrides& preferred_prefetch_overrides, - int64_t operand_size, const HloUse& hlo_use, - const absl::flat_hash_map& - instruction_schedule, - int64_t earliest_prefetch_time, int64_t latest_prefetch_time); - - static absl::StatusOr> - GetPrefetchOverrideInfo( - const MsaTensorOverrides& msa_tensor_overrides, - const PreferredPrefetchOverrides& preferred_prefetch_overrides, - int64_t operand_size, const HloUse& hlo_use, - const absl::flat_hash_map& - instruction_schedule, - int64_t earliest_prefetch_time, int64_t latest_prefetch_time, - const std::optional& position = std::nullopt); - - static bool ShouldPinInAlternateMemory( - const MsaTensorOverrides& msa_tensor_overrides, - const HloPosition& position, int64_t size); - - static bool ShouldKeepInDefaultMemory( - const MsaTensorOverrides& msa_tensor_overrides, - const HloPosition& position, int64_t size); - - static bool ShouldKeepInDefaultMemory( - const MsaTensorOverrides& msa_tensor_overrides, const HloUse& hlo_use, - int64_t operand_size); - - static bool DoesCrossProgramPrefetchBufferMatchAnyFilter( - const MsaSortOrderOverrides& sort_order_overrides, - const MsaBufferInterval& buffer_interval); - - // Returns an integer representing the priority of a MsaBufferInterval during - // assignment, a smaller number indicates a higher priority. - static int64_t GetBufferIntervalOverridePriority( - const MsaSortOrderOverrides& msa_sort_order_overrides, - const MsaBufferInterval& buffer_interval, - bool is_cross_program_prefetch = false); - - private: - static bool DoesBufferIntervalMatchHloUseFilter( - const HloPositionMatcher& filter, - const MsaBufferInterval& buffer_interval); -}; - -} // namespace memory_space_assignment -} // namespace xla - -#endif // XLA_SERVICE_MEMORY_SPACE_ASSIGNMENT_UTILS_H_ diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc b/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc index dc39566328f533..f1c1220d911487 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_stream.cc @@ -73,11 +73,8 @@ absl::Status LaunchSyclKernel( size_t num_args = num_args_ptr ? *num_args_ptr : 0; for (size_t arg_index = 0; arg_index < num_args; ++arg_index) { - if (arg_ptrs[arg_index] == nullptr) { - LOG(ERROR) << "LaunchSyclKernel: kernel argument " << arg_index - << " is null, cannot set kernel argument."; - return; - } + // A kernel argument can be null for a zero-sized buffer (e.g. empty + // scatter indices). VLOG(2) << "Setting kernel argument " << arg_index << " at address: " << arg_ptrs[arg_index]; cgh.set_arg(arg_index, arg_ptrs[arg_index]); diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index d0e406542294a3..eeebfd10a03b90 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -3674,8 +3674,10 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", + "@riegeli//riegeli/base:initializer", "@riegeli//riegeli/bytes:string_reader", "@riegeli//riegeli/bytes:string_writer", + "@tsl//tsl/platform:path", "@tsl//tsl/platform:protobuf", ], ) diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD index f023698959ff34..fa8e84382ea258 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD @@ -24,6 +24,7 @@ cc_library( "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@tsl//tsl/platform:path", ], ) @@ -41,12 +42,17 @@ xla_cc_test( "//xla/pjrt/proto:compile_options_proto_cc", "//xla/service/cpu:executable_proto_cc", "//xla/tests:aot_interception_pjrt_client", + "//xla/tsl/platform:env", "//xla/tsl/platform:test", "//xla/tsl/platform:test_main", "//xla/tsl/util/proto:parse_text_proto", "//xla/util/split_proto:human_readable_aot_executable_proto_cc", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", "@tsl//tsl/platform:path", + "@tsl//tsl/platform:protobuf", ], ) diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD index e458489ac816ae..944a9784e441ba 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD @@ -37,6 +37,7 @@ xla_test( deps = [ "//xla:literal", "//xla/service:hlo_module_config", + "//xla/tests:aot_interception_pjrt_client", "//xla/tests:literal_test_util", "//xla/tests:xla_internal_test_main", "//xla/tests/aot_compatibility_experimental:test_lib", diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc index 1c2b3c5d66e375..6dcdf613e6804e 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/service/hlo_module_config.h" #include "xla/tests/aot_compatibility_experimental/test_lib.h" +#include "xla/tests/aot_interception_pjrt_client.h" #include "xla/tests/literal_test_util.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/test.h" @@ -55,7 +56,7 @@ class CollectiveOpsAotTest : public AotCompatibilityTest, UnitTest::GetInstance()->current_test_info()->name(); std::string test_name = full_name.substr(0, full_name.find('/')); file_path_ = tsl::io::JoinPath( - GetExecutablesDirectory(GetParam().target_name), + GetExecutablesDirectory(GetParam().target_name, AOTTestPlatform::kGpu), absl::StrCat("v", version), absl::StrCat(test_name, ".pbtxt")); } }; @@ -141,7 +142,7 @@ std::vector GetTestParamsOrDie( INSTANTIATE_TEST_SUITE_P( BackwardsCompatibility, CollectiveOpsAotTest, ValuesIn(GetTestParamsOrDie(GetAotTestParamsForBackwardsCompatibility( - "collective_ops_aot_test_2gpu"))), + "collective_ops_aot_test_2gpu", AOTTestPlatform::kGpu))), [](const TestParamInfo& info) { return absl::StrCat("v", info.param.version); }); @@ -149,7 +150,7 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P( GoldenFileVerification, CollectiveOpsAotTest, ValuesIn(GetTestParamsOrDie(GetAotTestParamsForGoldenFileVerification( - "collective_ops_aot_test_2gpu"))), + "collective_ops_aot_test_2gpu", AOTTestPlatform::kGpu))), [](const TestParamInfo& info) { return absl::StrCat("v", info.param.version); }); diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/test_dummy_test/v1/exec.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/test_dummy_test/v1/exec.pbtxt index e69de29bb2d1d6..23cca5efb740ec 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/test_dummy_test/v1/exec.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/test_dummy_test/v1/exec.pbtxt @@ -0,0 +1,7 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable + +# Minimal well-formed GPU artifact +gpu_executable { + binary: "dummy_gpu_binary" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc index d11757ce29d755..98d1c661a30001 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc @@ -24,6 +24,7 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" @@ -43,11 +44,49 @@ namespace aot_compatibility_experimental { using ::testing::TestInfo; using ::testing::UnitTest; +namespace { + +absl::StatusOr> GetExecutableVersions( + absl::string_view target_name, AOTTestPlatform platform) { + std::string dir = GetExecutablesDirectory(target_name, platform); + ABSL_ASSIGN_OR_RETURN(std::vector versions, + test_lib_internal::GetExecutableVersionsInDir(dir)); + // On some filesystems a missing directory lists empty rather than failing, so + // emptiness -- not the status -- is what identifies a bad or ungenerated + // target. + if (versions.empty()) { + return absl::NotFoundError( + absl::StrCat("No AOT golden artifacts found for target '", target_name, + "' under ", dir)); + } + return versions; +} + +// Treats an environment variable as a boolean flag: `=0` must disable golden +// dumping, not enable it. An unparseable value counts as disabled. +bool IsEnvFlagEnabled(const char* name) { + const char* value = std::getenv(name); + bool enabled = false; + return value != nullptr && absl::SimpleAtob(value, &enabled) && enabled; +} + +// The undeclared-outputs directory exists only when the golden-update helper +// drives the test. Shared so both callers enforce it with the same message. +absl::StatusOr GetUndeclaredOutputsDir() { + const char* out_dir = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); + if (out_dir == nullptr) { + return absl::FailedPreconditionError( + "TEST_UNDECLARED_OUTPUTS_DIR is unset. Golden updates must be driven " + "by the golden-update helper for this package, which sets it; running " + "the test target directly will not work."); + } + return std::string(out_dir); +} + +} // namespace + std::string GetExecutablesDirectory(absl::string_view target_name, AOTTestPlatform platform) { - // We use the full target name as part of the path, including backend (e.g. - // collective_ops_aot_test_2gpu). The platform selects the "gpu" or "cpu" - // executables subdirectory via the shared PlatformSubdir mapping. return tsl::io::JoinPath( tsl::testing::TensorFlowSrcRoot(), absl::StrCat("compiler/xla/tests/aot_compatibility_experimental/", @@ -56,52 +95,53 @@ std::string GetExecutablesDirectory(absl::string_view target_name, target_name); } -namespace { +namespace test_lib_internal { -// Returns all available artifact versions sorted in ascending order. -absl::StatusOr> GetExecutableVersions( - absl::string_view target_name, AOTTestPlatform platform) { - std::string dir = GetExecutablesDirectory(target_name, platform); +absl::StatusOr> GetExecutableVersionsInDir( + absl::string_view dir) { std::vector children; auto* env = tsl::Env::Default(); - ABSL_RETURN_IF_ERROR(env->GetChildren(dir, &children)); + ABSL_RETURN_IF_ERROR(env->GetChildren(std::string(dir), &children)); - std::vector all_versions; - all_versions.reserve(children.size()); + std::vector versions; + versions.reserve(children.size()); for (const std::string& child : children) { - if (!absl::StartsWith(child, "v")) { - return absl::InvalidArgumentError( - absl::StrCat("Failed to parse version: ", child)); - } - std::string child_path = tsl::io::JoinPath(dir, child); - ABSL_RETURN_IF_ERROR(env->IsDirectory(child_path)); - absl::string_view version_str = absl::string_view(child).substr(1); + // Skip anything that is not a `v` directory holding a golden; a stray or + // empty entry must not wedge test discovery for the whole target. int32_t version; - if (!absl::SimpleAtoi(version_str, &version)) { - return absl::InvalidArgumentError( - absl::StrCat("Failed to parse version: ", child)); + if (!absl::StartsWith(child, "v") || + !absl::SimpleAtoi(absl::string_view(child).substr(1), &version)) { + continue; } - all_versions.push_back(version); + std::vector files; + if (!env->GetChildren(tsl::io::JoinPath(dir, child), &files).ok() || + std::none_of(files.begin(), files.end(), [](absl::string_view f) { + return absl::EndsWith(f, ".pbtxt"); + })) { + continue; + } + versions.push_back(version); } - std::sort(all_versions.begin(), all_versions.end()); - return all_versions; + std::sort(versions.begin(), versions.end()); + return versions; } -} // namespace +} // namespace test_lib_internal absl::StatusOr> GetAotTestParamsForBackwardsCompatibility(absl::string_view target_name, AOTTestPlatform platform) { + if (IsEnvFlagEnabled("XLA_AOT_UPDATE_GOLDENS")) { + return std::vector{}; + } ABSL_ASSIGN_OR_RETURN(std::vector versions, GetExecutableVersions(target_name, platform)); - if (std::getenv("XLA_AOT_TEST_ALL_VERSIONS") == nullptr && - versions.size() > 2) { - // For backwards compatibility testing, we only test the minimum and the - // (maximum - 1) versions to verify the boundaries of our compatibility - // guarantees. The maximum version is omitted here because it is already - // covered by the golden file verification. + if (!IsEnvFlagEnabled("XLA_AOT_TEST_ALL_VERSIONS") && versions.size() > 2) { + // The oldest and second-newest versions bound our compatibility guarantee. + // This is positional, not arithmetic: for v1, v2, v5 the pair is {v1, v2}. + // The newest is covered by golden file verification instead. versions = {versions.front(), versions[versions.size() - 2]}; } @@ -117,12 +157,13 @@ GetAotTestParamsForBackwardsCompatibility(absl::string_view target_name, absl::StatusOr> GetAotTestParamsForGoldenFileVerification(absl::string_view target_name, AOTTestPlatform platform) { + if (IsEnvFlagEnabled("XLA_AOT_UPDATE_GOLDENS")) { + ABSL_RETURN_IF_ERROR(GetUndeclaredOutputsDir().status()); + return std::vector{ + {AOTTestMode::kUpdateGolden, 0, std::string(target_name)}}; + } ABSL_ASSIGN_OR_RETURN(std::vector versions, GetExecutableVersions(target_name, platform)); - if (versions.empty()) { - return absl::NotFoundError( - absl::StrCat("No artifacts found for target: ", target_name)); - } return std::vector{{AOTTestMode::kGoldenVerification, versions.back(), std::string(target_name)}}; @@ -143,20 +184,23 @@ AotCompatibilityTest::AotCompatibilityTest(AotTestParam param) << (*client)->platform_name() << ". " << platform.status(); const TestInfo* test_info = UnitTest::GetInstance()->current_test_info(); - std::string test_name = ""; - if (test_info != nullptr) { - absl::string_view name_view = test_info->name(); - size_t slash_pos = name_view.find('/'); - if (slash_pos != absl::string_view::npos) { - test_name = std::string(name_view.substr(0, slash_pos)); - } else { - test_name = std::string(name_view); - } + // gtest names a parameterized case `/`; the + // golden is named after the test alone. + absl::string_view name = + test_info == nullptr ? "" : test_info->name(); + std::string test_name(name.substr(0, name.find('/'))); + std::string artifact_path; + if (param.mode == AOTTestMode::kUpdateGolden) { + absl::StatusOr out_dir = GetUndeclaredOutputsDir(); + CHECK_OK(out_dir.status()); + artifact_path = tsl::io::JoinPath( + *out_dir, absl::StrCat(test_name, ".pbtxt")); + } else { + artifact_path = tsl::io::JoinPath( + GetExecutablesDirectory(param.target_name, platform.value()), + absl::StrCat("v", param.version), + absl::StrCat(test_name, ".pbtxt")); } - std::string artifact_path = tsl::io::JoinPath( - GetExecutablesDirectory(param.target_name, platform.value()), - absl::StrCat("v", param.version), - absl::StrCat(test_name, ".pbtxt")); return std::make_unique( std::move(*client), param.mode, artifact_path); }(param), @@ -164,7 +208,7 @@ AotCompatibilityTest::AotCompatibilityTest(AotTestParam param) DebugOptions AotCompatibilityTest::GetDebugOptionsForTest() const { DebugOptions debug_options = HloTestBase::GetDebugOptionsForTest(); - debug_options.set_xla_gpu_exclude_nondeterministic_ops(true); + AOTInterceptionPjrtClient::ApplyAotDeterminismDebugOptions(debug_options); return debug_options; } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h index 71302210ecb241..5f0b89b9bd71bb 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h @@ -20,6 +20,7 @@ limitations under the License. #include #include "absl/status/statusor.h" +#include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "xla/tests/aot_interception_pjrt_client.h" #include "xla/tests/hlo_test_base.h" @@ -37,30 +38,42 @@ struct AotTestParam { return mode == other.mode && version == other.version && target_name == other.target_name; } + + // Without this gtest prints the raw object bytes on a failed match. + template + friend void AbslStringify(Sink& sink, const AotTestParam& param) { + absl::Format(&sink, "AotTestParam{mode=%d, version=%d, target_name=%s}", + static_cast(param.mode), param.version, + param.target_name); + } }; -// Returns the path to the executables directory for the current test target and -// platform (defaults to GPU for backwards compatibility with existing callers). -std::string GetExecutablesDirectory( - absl::string_view target_name, - AOTTestPlatform platform = AOTTestPlatform::kGpu); - -// Gets the list of AOT test parameters for testing backwards compatibility -// boundaries. -// By default we test only 2 versions for backwards compatibility: the minimum -// and the (maximum - 1) versions to verify the boundaries of our compatibility -// guarantees. Set XLA_AOT_TEST_ALL_VERSIONS to test all versions. +// Returns the path to the executables directory for the given test target and +// platform. +std::string GetExecutablesDirectory(absl::string_view target_name, + AOTTestPlatform platform); + +// Returns the parameters bounding our backwards compatibility guarantee: the +// oldest and second-newest versions. Set XLA_AOT_TEST_ALL_VERSIONS=1 to +// return all available versions. absl::StatusOr> -GetAotTestParamsForBackwardsCompatibility( - absl::string_view target_name, - AOTTestPlatform platform = AOTTestPlatform::kGpu); +GetAotTestParamsForBackwardsCompatibility(absl::string_view target_name, + AOTTestPlatform platform); -// Returns the latest version of the AOT dumped artifact, wrapped in a list for -// test parameterization. +// Returns the latest version, wrapped in a list for test parameterization. absl::StatusOr> -GetAotTestParamsForGoldenFileVerification( - absl::string_view target_name, - AOTTestPlatform platform = AOTTestPlatform::kGpu); +GetAotTestParamsForGoldenFileVerification(absl::string_view target_name, + AOTTestPlatform platform); + +// Implementation details, exposed only so test_lib_test can exercise them. +namespace test_lib_internal { + +// Returns, in ascending numeric order, the version numbers of the `v` +// subdirectories of `dir` that hold at least one .pbtxt. Others are skipped. +absl::StatusOr> GetExecutableVersionsInDir( + absl::string_view dir); + +} // namespace test_lib_internal // A parameterized test fixture base class for AOT compatibility tests. class AotCompatibilityTest : public HloTestBase { diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib_test.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib_test.cc index 1aa91819b4d225..ffb987c944b9ce 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib_test.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib_test.cc @@ -17,18 +17,25 @@ limitations under the License. #include +#include +#include #include #include +#include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_matchers.h" +#include "absl/strings/match.h" +#include "absl/strings/string_view.h" #include "xla/pjrt/proto/compile_options.pb.h" #include "xla/service/cpu/executable.pb.h" #include "xla/tests/aot_interception_pjrt_client.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/test.h" #include "xla/tsl/util/proto/parse_text_proto.h" #include "xla/util/split_proto/human_readable_aot_executable.pb.h" #include "tsl/platform/path.h" +#include "tsl/platform/protobuf.h" namespace xla { namespace aot_compatibility_experimental { @@ -38,13 +45,82 @@ using ::absl_testing::StatusIs; using ::testing::AllOf; using ::testing::ElementsAre; using ::testing::HasSubstr; +using ::testing::IsEmpty; +using ::testing::Not; +using ::testing::TempDir; using ::tsl::proto_testing::ParseTextProtoOrDie; +// Sets an environment variable for the lifetime of the object and restores the +// previous state -- including absence -- afterwards. Every case in this file +// runs in one process, so an unrestored variable silently changes the meaning +// of every case that follows it. +class ScopedEnv { + public: + // A null `value` unsets the variable instead of assigning to it. + ScopedEnv(const char* name, const char* value) : name_(name) { + if (const char* previous = std::getenv(name)) { + had_previous_ = true; + previous_ = previous; + } + if (value == nullptr) { + unsetenv(name); + } else { + setenv(name, value, /*overwrite=*/1); + } + } + + ScopedEnv(const ScopedEnv&) = delete; + ScopedEnv& operator=(const ScopedEnv&) = delete; + + ~ScopedEnv() { + if (had_previous_) { + setenv(name_, previous_.c_str(), /*overwrite=*/1); + } else { + unsetenv(name_); + } + } + + private: + const char* name_; + bool had_previous_ = false; + std::string previous_; +}; + +// A temporary directory private to one test case. `TempDir()` is +// shared by the whole binary, so a case that scans it can otherwise observe +// entries another case created. +class ScopedTempDir { + public: + explicit ScopedTempDir(absl::string_view name) + : path_(tsl::io::JoinPath(TempDir(), name)) { + Remove(); + CHECK_OK(tsl::Env::Default()->RecursivelyCreateDir(path_)); + } + + ScopedTempDir(const ScopedTempDir&) = delete; + ScopedTempDir& operator=(const ScopedTempDir&) = delete; + + ~ScopedTempDir() { Remove(); } + + const std::string& path() const { return path_; } + + private: + void Remove() { + int64_t undeleted_files = 0; + int64_t undeleted_dirs = 0; + tsl::Env::Default() + ->DeleteRecursively(path_, &undeleted_files, &undeleted_dirs) + .IgnoreError(); + } + + std::string path_; +}; + TEST(TestLibTest, GetAotTestParamsForBackwardsCompatibility_With4Versions) { - unsetenv("XLA_AOT_TEST_ALL_VERSIONS"); - ASSERT_OK_AND_ASSIGN( - std::vector params, - GetAotTestParamsForBackwardsCompatibility("test_dummy_test")); + ScopedEnv all_versions("XLA_AOT_TEST_ALL_VERSIONS", nullptr); + ASSERT_OK_AND_ASSIGN(std::vector params, + GetAotTestParamsForBackwardsCompatibility( + "test_dummy_test", AOTTestPlatform::kGpu)); EXPECT_THAT(params, ElementsAre(AotTestParam{AOTTestMode::kBackwardsCompatibility, 1, "test_dummy_test"}, @@ -53,10 +129,10 @@ TEST(TestLibTest, GetAotTestParamsForBackwardsCompatibility_With4Versions) { } TEST(TestLibTest, GetAotTestParamsForBackwardsCompatibility_AllVersions) { - setenv("XLA_AOT_TEST_ALL_VERSIONS", "1", 1); - ASSERT_OK_AND_ASSIGN( - std::vector params, - GetAotTestParamsForBackwardsCompatibility("test_dummy_test")); + ScopedEnv all_versions("XLA_AOT_TEST_ALL_VERSIONS", "1"); + ASSERT_OK_AND_ASSIGN(std::vector params, + GetAotTestParamsForBackwardsCompatibility( + "test_dummy_test", AOTTestPlatform::kGpu)); EXPECT_THAT(params, ElementsAre(AotTestParam{AOTTestMode::kBackwardsCompatibility, 1, "test_dummy_test"}, @@ -66,13 +142,85 @@ TEST(TestLibTest, GetAotTestParamsForBackwardsCompatibility_AllVersions) { "test_dummy_test"}, AotTestParam{AOTTestMode::kBackwardsCompatibility, 4, "test_dummy_test"})); - unsetenv("XLA_AOT_TEST_ALL_VERSIONS"); +} + +// The kUpdateGolden control plane. Under XLA_AOT_UPDATE_GOLDENS the two +// parameter generators change behaviour completely, and none of it needs +// hardware or a PjRt client. + +TEST(TestLibTest, + GetAotTestParamsForBackwardsCompatibilityUpdateGoldensIsEmpty) { + ScopedEnv update_goldens("XLA_AOT_UPDATE_GOLDENS", "1"); + // Regenerating goldens recompiles; replaying old executables is meaningless + // in that mode, so the suite must be instantiated with no parameters at all. + ASSERT_OK_AND_ASSIGN(std::vector params, + GetAotTestParamsForBackwardsCompatibility( + "test_dummy_test", AOTTestPlatform::kGpu)); + EXPECT_THAT(params, IsEmpty()); +} + +TEST(TestLibTest, + GetAotTestParamsForGoldenFileVerificationUpdateGoldensYieldsOneParam) { + ScopedEnv update_goldens("XLA_AOT_UPDATE_GOLDENS", "1"); + const std::string outputs_path = TempDir(); + ScopedEnv outputs_dir("TEST_UNDECLARED_OUTPUTS_DIR", outputs_path.c_str()); + // Version 0 is the sentinel: the dump does not read an existing version. + ASSERT_OK_AND_ASSIGN(std::vector params, + GetAotTestParamsForGoldenFileVerification( + "test_dummy_test", AOTTestPlatform::kGpu)); + EXPECT_THAT(params, ElementsAre(AotTestParam{AOTTestMode::kUpdateGolden, 0, + "test_dummy_test"})); +} + +TEST(TestLibTest, + GetAotTestParamsForGoldenFileVerificationUpdateGoldensNeedsOutputsDir) { + ScopedEnv update_goldens("XLA_AOT_UPDATE_GOLDENS", "1"); + ScopedEnv outputs_dir("TEST_UNDECLARED_OUTPUTS_DIR", nullptr); + // Without the outputs directory the dump has nowhere to write, which means + // the test was invoked directly rather than through update_goldens.sh. + EXPECT_THAT(GetAotTestParamsForGoldenFileVerification("test_dummy_test", + AOTTestPlatform::kGpu), + StatusIs(absl::StatusCode::kFailedPrecondition, + HasSubstr("TEST_UNDECLARED_OUTPUTS_DIR"))); +} + +TEST(TestLibTest, + GetAotTestParamsForBackwardsCompatibilityMissingTargetDirIsNotFound) { + ScopedEnv all_versions("XLA_AOT_TEST_ALL_VERSIONS", nullptr); + ScopedEnv update_goldens("XLA_AOT_UPDATE_GOLDENS", nullptr); + // A target whose executables directory does not exist is a mistake -- a + // typo'd name, or goldens that were never generated. Returning no parameters + // would leave the suite green while testing nothing, so it must be an error, + // and the same error the verification generator produces. + EXPECT_THAT(GetAotTestParamsForBackwardsCompatibility( + "target_that_does_not_exist", AOTTestPlatform::kGpu), + StatusIs(absl::StatusCode::kNotFound)); +} + +TEST(TestLibTest, + GetAotTestParamsForGoldenFileVerificationMissingTargetDirIsNotFound) { + ScopedEnv update_goldens("XLA_AOT_UPDATE_GOLDENS", nullptr); + EXPECT_THAT(GetAotTestParamsForGoldenFileVerification( + "target_that_does_not_exist", AOTTestPlatform::kGpu), + StatusIs(absl::StatusCode::kNotFound)); +} + +TEST(TestLibTest, UpdateGoldensEnvFlagIsTruthyNotMerelyPresent) { + ScopedEnv all_versions("XLA_AOT_TEST_ALL_VERSIONS", nullptr); + // `=0` must not enable golden dumping. Presence-only checks made + // `--test_env=XLA_AOT_UPDATE_GOLDENS=0` turn dumping on, which is the + // opposite of what anyone passing it intends. + ScopedEnv update_goldens("XLA_AOT_UPDATE_GOLDENS", "0"); + ASSERT_OK_AND_ASSIGN(std::vector params, + GetAotTestParamsForBackwardsCompatibility( + "test_dummy_test", AOTTestPlatform::kGpu)); + EXPECT_THAT(params, Not(IsEmpty())); } TEST(TestLibTest, GetAotTestParamsForGoldenFileVerification_With4Versions) { - ASSERT_OK_AND_ASSIGN( - std::vector params, - GetAotTestParamsForGoldenFileVerification("test_dummy_test")); + ASSERT_OK_AND_ASSIGN(std::vector params, + GetAotTestParamsForGoldenFileVerification( + "test_dummy_test", AOTTestPlatform::kGpu)); EXPECT_THAT(params, ElementsAre(AotTestParam{AOTTestMode::kGoldenVerification, 4, "test_dummy_test"})); } @@ -80,7 +228,8 @@ TEST(TestLibTest, GetAotTestParamsForGoldenFileVerification_With4Versions) { TEST(TestLibTest, AOTInterceptionPjrtClientPackArtifactForInnerClient_Succeeds) { std::string artifact_path = tsl::io::JoinPath( - GetExecutablesDirectory("test_dummy_test"), "v1", "exec.pbtxt"); + GetExecutablesDirectory("test_dummy_test", AOTTestPlatform::kGpu), "v1", + "exec.pbtxt"); AOTInterceptionPjrtClient client( nullptr, AOTTestMode::kBackwardsCompatibility, artifact_path); ASSERT_OK_AND_ASSIGN(std::string serialized, @@ -263,6 +412,103 @@ TEST(TestLibTest, DeserializeToHumanReadable_CpuUnpacksPlainProto) { unpacked.executable_and_options().serialized_executable().empty()); } +TEST(TestLibTest, GetExecutableVersionsIgnoresEmptyDirs) { + ScopedTempDir temp_dir("ignores_empty_dirs"); + tsl::Env* env = tsl::Env::Default(); + + // v1 contains a golden. + std::string v1_dir = tsl::io::JoinPath(temp_dir.path(), "v1"); + ASSERT_OK(env->RecursivelyCreateDir(v1_dir)); + ASSERT_OK( + tsl::WriteStringToFile(env, tsl::io::JoinPath(v1_dir, "test.pbtxt"), "")); + + // v2 is an empty leftover, e.g. from syncing to a revision without goldens. + std::string v2_dir = tsl::io::JoinPath(temp_dir.path(), "v2"); + ASSERT_OK(env->RecursivelyCreateDir(v2_dir)); + + EXPECT_THAT(test_lib_internal::GetExecutableVersionsInDir(temp_dir.path()), + absl_testing::IsOkAndHolds(ElementsAre(1))); +} + +// A `v`-prefixed regular file and a non-numeric `v` directory are unrelated +// entries that must not wedge test discovery for the whole target: the scanner +// skips them and still reports the real versions. +TEST(TestLibTest, GetExecutableVersionsSkipsNonVersionEntries) { + ScopedTempDir temp_dir("skips_non_version_entries"); + tsl::Env* env = tsl::Env::Default(); + + std::string v1_dir = tsl::io::JoinPath(temp_dir.path(), "v1"); + ASSERT_OK(env->RecursivelyCreateDir(v1_dir)); + ASSERT_OK( + tsl::WriteStringToFile(env, tsl::io::JoinPath(v1_dir, "test.pbtxt"), "")); + + // A `v`-prefixed regular file, not a directory. + ASSERT_OK(tsl::WriteStringToFile( + env, tsl::io::JoinPath(temp_dir.path(), "v9.txt"), "")); + + // A `v`-prefixed directory whose suffix is not a number. It holds a .pbtxt + // so it survives the empty-directory filter and genuinely reaches the parse. + std::string v_old_dir = tsl::io::JoinPath(temp_dir.path(), "v_old"); + ASSERT_OK(env->RecursivelyCreateDir(v_old_dir)); + ASSERT_OK(tsl::WriteStringToFile( + env, tsl::io::JoinPath(v_old_dir, "stale.pbtxt"), "")); + + // Entries that do not start with `v` at all: a plain file and a directory + // holding a .pbtxt, either of which could appear beside the goldens. + ASSERT_OK(tsl::WriteStringToFile( + env, tsl::io::JoinPath(temp_dir.path(), "README.md"), "")); + std::string notes_dir = tsl::io::JoinPath(temp_dir.path(), "notes"); + ASSERT_OK(env->RecursivelyCreateDir(notes_dir)); + ASSERT_OK( + tsl::WriteStringToFile(env, tsl::io::JoinPath(notes_dir, "n.pbtxt"), "")); + + EXPECT_THAT(test_lib_internal::GetExecutableVersionsInDir(temp_dir.path()), + absl_testing::IsOkAndHolds(ElementsAre(1))); +} + +// Versions must be ordered numerically, not lexically. With only v1..v4 on +// disk the two orderings agree, so this uses v10 to tell them apart: lexically +// "v10" sorts before "v2". +TEST(TestLibTest, GetExecutableVersionsSortsNumericallyNotLexically) { + ScopedTempDir temp_dir("sorts_numerically"); + tsl::Env* env = tsl::Env::Default(); + + for (absl::string_view name : {"v2", "v10", "v1"}) { + std::string dir = tsl::io::JoinPath(temp_dir.path(), name); + ASSERT_OK(env->RecursivelyCreateDir(dir)); + ASSERT_OK( + tsl::WriteStringToFile(env, tsl::io::JoinPath(dir, "test.pbtxt"), "")); + } + + EXPECT_THAT(test_lib_internal::GetExecutableVersionsInDir(temp_dir.path()), + absl_testing::IsOkAndHolds(ElementsAre(1, 2, 10))); +} + +TEST(TestLibTest, WriteGoldenTextProtoRoundTripWithHeader) { + ScopedTempDir temp_dir("write_golden_round_trip"); + std::string filepath = tsl::io::JoinPath(temp_dir.path(), "roundtrip.pbtxt"); + + HumanReadableAotExecutable original; + original.mutable_cpu_executable()->set_entry_function_name("roundtrip_test"); + + ASSERT_OK( + AOTInterceptionPjrtClient::WriteGoldenTextProto(original, filepath)); + + std::string file_content; + ASSERT_OK( + tsl::ReadFileToString(tsl::Env::Default(), filepath, &file_content)); + + EXPECT_TRUE(absl::StartsWith(file_content, "# proto-file: ")); + EXPECT_THAT(file_content, + HasSubstr("# proto-message: xla.HumanReadableAotExecutable\n")); + + HumanReadableAotExecutable parsed; + ASSERT_TRUE( + tsl::protobuf::TextFormat::ParseFromString(file_content, &parsed)); + EXPECT_EQ(parsed.cpu_executable().entry_function_name(), + original.cpu_executable().entry_function_name()); +} + } // namespace } // namespace aot_compatibility_experimental } // namespace xla diff --git a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc index 52d87c03958f23..b308f43752c725 100644 --- a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc +++ b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "riegeli/base/maker.h" #include "riegeli/bytes/string_reader.h" #include "riegeli/bytes/string_writer.h" #include "xla/hlo/builder/xla_computation.h" @@ -52,16 +53,16 @@ limitations under the License. #include "xla/util/split_proto/split_proto_reader.h" #include "xla/xla.pb.h" #include "xla/xla_data.pb.h" +#include "tsl/platform/path.h" #include "tsl/platform/protobuf.h" namespace xla { namespace { -// Resolves a nested field descriptor by walking `path` (a sequence of field -// names) starting from `root`. Every element except the last must name a -// message field. CHECK-fails if the schema no longer matches so that a stale -// path is caught immediately instead of silently ignoring nothing. +// Resolves a nested field descriptor by walking `path` from `root`; every +// element but the last must name a message field. CHECK-fails on a stale path +// rather than silently ignoring nothing. const tsl::protobuf::FieldDescriptor* FieldByPath( const tsl::protobuf::Descriptor* root, std::initializer_list path) { @@ -82,8 +83,7 @@ const tsl::protobuf::FieldDescriptor* FieldByPath( } // Runs the structural comparison shared by all backends. `extra_ignored_fields` -// lists additional field descriptors to ignore, on top of the common, -// backend-agnostic ones. +// adds to the common, backend-agnostic ignore list. absl::Status CompareStructurally( const HumanReadableAotExecutable& fresh, const HumanReadableAotExecutable& golden, @@ -126,41 +126,40 @@ AOTInterceptionPjrtClient::DeserializeToHumanReadable( absl::string_view serialized, AOTTestPlatform platform) { HumanReadableAotExecutable unpacked; - if (platform == AOTTestPlatform::kGpu) { - ABSL_RETURN_IF_ERROR( - ReadSplitProto(std::make_unique>(serialized), - *unpacked.mutable_executable_and_options())); - - std::string serialized_gpu_exec = - std::move(*unpacked.mutable_executable_and_options() - ->mutable_serialized_executable()); - unpacked.mutable_executable_and_options()->clear_serialized_executable(); - - ABSL_RETURN_IF_ERROR( - ReadSplitProto(std::make_unique>( - std::move(serialized_gpu_exec)), - *unpacked.mutable_gpu_executable())); - return unpacked; - } - - // CPU artifacts are plain (non-split) protos: an ExecutableAndOptionsProto - // whose `serialized_executable` field holds a serialized - // cpu::CompilationResultProto. - ExecutableAndOptionsProto& executable_and_options = - *unpacked.mutable_executable_and_options(); - if (!executable_and_options.ParseFromString(serialized)) { - return absl::InternalError( - "AOTInterceptionPjrtClient: failed to parse CPU " - "ExecutableAndOptionsProto."); - } - if (!unpacked.mutable_cpu_executable()->ParseFromString( - executable_and_options.serialized_executable())) { - return absl::InternalError( - "AOTInterceptionPjrtClient: failed to parse CPU " - "CompilationResultProto."); + switch (platform) { + case AOTTestPlatform::kGpu: { + ABSL_RETURN_IF_ERROR( + ReadSplitProto(std::make_unique>(serialized), + *unpacked.mutable_executable_and_options())); + + ABSL_RETURN_IF_ERROR(ReadSplitProto( + riegeli::Maker( + unpacked.executable_and_options().serialized_executable()), + *unpacked.mutable_gpu_executable())); + unpacked.mutable_executable_and_options()->clear_serialized_executable(); + return unpacked; + } + case AOTTestPlatform::kCpu: { + // CPU artifacts are plain protos: an ExecutableAndOptionsProto whose + // `serialized_executable` holds a cpu::CompilationResultProto. + ExecutableAndOptionsProto& executable_and_options = + *unpacked.mutable_executable_and_options(); + if (!executable_and_options.ParseFromString(serialized)) { + return absl::InternalError( + "AOTInterceptionPjrtClient: failed to parse CPU " + "ExecutableAndOptionsProto."); + } + if (!unpacked.mutable_cpu_executable()->ParseFromString( + executable_and_options.serialized_executable())) { + return absl::InternalError( + "AOTInterceptionPjrtClient: failed to parse CPU " + "CompilationResultProto."); + } + executable_and_options.clear_serialized_executable(); + return unpacked; + } } - executable_and_options.clear_serialized_executable(); - return unpacked; + CHECK(false) << "Unsupported platform for AOT testing!"; } absl::StatusOr AOTInterceptionPjrtClient::PlatformFromName( @@ -183,23 +182,19 @@ absl::string_view AOTInterceptionPjrtClient::PlatformSubdir( return "gpu"; case AOTTestPlatform::kCpu: return "cpu"; - default: - CHECK(false) << "Unsupported platform for AOT testing!"; } + CHECK(false) << "Unsupported platform for AOT testing!"; } absl::Status AOTInterceptionPjrtClient::CompareGPUExecutables( const HumanReadableAotExecutable& fresh, const HumanReadableAotExecutable& golden) { - // `binary`, `asm_text`, `buffer_allocations`, `gpu_compute_capability` and - // the `ptx`/`cubin` kernel binaries capture backend machine code and - // device-specific details and are not part of the structural comparison. + // The ignored fields hold backend machine code and device-specific details. // - // TODO(b/528258781): Debug options are currently ignored wholesale (both the - // shared build options and the HLO module config). Evaluate which flags - // meaningfully affect the compiled artifact and should be compared, versus - // which are host- or run-specific noise (e.g. dump paths and cache - // directories) that must be excluded, and ignore only the latter. + // TODO(b/528258781): Debug options are ignored wholesale. Work out which + // flags actually affect the artifact and should be compared, versus which are + // host- or run-specific noise (dump paths, cache dirs), and ignore only + // those. return CompareStructurally( fresh, golden, { @@ -209,10 +204,8 @@ absl::Status AOTInterceptionPjrtClient::CompareGPUExecutables( "buffer_allocations"), gpu::GpuExecutableProto::descriptor()->FindFieldByName( "gpu_compute_capability"), - // `ptx`/`cubin` hold the (non-deterministic) kernel binaries embedded - // in custom-kernel thunks; excluded like `binary`/`asm_text`. Both - // oneof variants are ignored, so a ptx<->cubin case flip is - // intentionally tolerated (same kernel, different backend encoding). + // Non-deterministic kernel binaries in custom-kernel thunks. Both + // oneof variants are ignored, so a ptx<->cubin flip is tolerated. stream_executor::KernelLoaderSpecProto::descriptor()->FindFieldByName( "ptx"), stream_executor::KernelLoaderSpecProto::descriptor()->FindFieldByName( @@ -226,10 +219,8 @@ absl::Status AOTInterceptionPjrtClient::CompareGPUExecutables( absl::Status AOTInterceptionPjrtClient::CompareGoldenCPUExecutable( const HumanReadableAotExecutable& fresh, const HumanReadableAotExecutable& golden) { - // `object_files`, `target_machine_options` and `data_layout` capture the - // compiled machine code and host-specific target details; `debug_options` - // (both the shared build options and the HLO module config) is host- or - // run-specific noise. None are part of the structural comparison. + // The ignored fields hold compiled machine code, host-specific target details + // and debug options, which are host- or run-specific noise. return CompareStructurally( fresh, golden, { @@ -242,10 +233,8 @@ absl::Status AOTInterceptionPjrtClient::CompareGoldenCPUExecutable( ExecutableBuildOptionsProto::descriptor()->FindFieldByName( "debug_options"), HloModuleConfigProto::descriptor()->FindFieldByName("debug_options"), - // The following capture host/process-specific metadata (a global - // module-id counter and the host intra-op thread-pool size) that - // legitimately varies between the golden-generation run and the test - // run, so they are excluded from the structural comparison. + // Host/process metadata: a global module-id counter and the intra-op + // thread-pool size, both of which vary run to run. FieldByPath(cpu::CompilationResultProto::descriptor(), {"hlo_module", "hlo_module", "id"}), FieldByPath(cpu::CompilationResultProto::descriptor(), @@ -267,15 +256,75 @@ absl::Status AOTInterceptionPjrtClient::VerifyAgainstGolden( ABSL_ASSIGN_OR_RETURN(HumanReadableAotExecutable golden, LoadHumanReadableArtifact()); + absl::Status diff_status; switch (platform) { case AOTTestPlatform::kGpu: - return CompareGPUExecutables(fresh_unpacked, golden); + diff_status = CompareGPUExecutables(fresh_unpacked, golden); + break; case AOTTestPlatform::kCpu: - return CompareGoldenCPUExecutable(fresh_unpacked, golden); + diff_status = CompareGoldenCPUExecutable(fresh_unpacked, golden); + break; } - return absl::UnimplementedError( - "AOTInterceptionPjrtClient: unsupported platform in " - "VerifyAgainstGolden."); + + if (!diff_status.ok()) { + absl::string_view target_name = + tsl::io::Basename(tsl::io::Dirname(tsl::io::Dirname(artifact_path_))); + return absl::Status( + diff_status.code(), + absl::StrCat( + "Golden artifact verification failed for ", artifact_path_, ".\n", + diff_status.message(), "\n\nRegenerate the goldens for target `", + target_name, + "` and add the newly generated v directory to the change " + "after reviewing the differences.")); + } + + return absl::OkStatus(); +} + +void AOTInterceptionPjrtClient::ApplyAotDeterminismDebugOptions( + DebugOptions& debug_options) { + debug_options.set_xla_gpu_exclude_nondeterministic_ops(true); +} + +absl::Status AOTInterceptionPjrtClient::DumpGolden( + const PjRtExecutable& fresh_executable) { + ABSL_ASSIGN_OR_RETURN(std::string serialized, + fresh_executable.SerializeExecutable()); + ABSL_ASSIGN_OR_RETURN(AOTTestPlatform platform, + PlatformFromName(inner_client_->platform_name())); + ABSL_ASSIGN_OR_RETURN(HumanReadableAotExecutable unpacked, + DeserializeToHumanReadable(serialized, platform)); + return WriteGoldenTextProto(unpacked, artifact_path_); +} + +absl::Status AOTInterceptionPjrtClient::WriteGoldenTextProto( + const HumanReadableAotExecutable& unpacked, absl::string_view path) { + tsl::Env* env = tsl::Env::Default(); + ABSL_RETURN_IF_ERROR( + env->RecursivelyCreateDir(std::string(tsl::io::Dirname(path)))); + // The text-format schema header lets editors and formatters recognize it. + const tsl::protobuf::Descriptor* descriptor = unpacked.GetDescriptor(); + std::string body; + if (!tsl::protobuf::TextFormat::PrintToString(unpacked, &body)) { + return absl::InternalError( + absl::StrCat("Failed to serialize HumanReadableAotExecutable to text " + "format for ", + path)); + } + std::string out = + absl::StrCat("# proto-file: ", descriptor->file()->name(), "\n", + "# proto-message: ", descriptor->full_name(), "\n", body); + // Write to a sibling temporary and rename, so an interrupted or failed write + // cannot leave a truncated golden behind for someone to check in. + const std::string tmp_path = absl::StrCat(path, ".tmp"); + ABSL_RETURN_IF_ERROR(tsl::WriteStringToFile(env, tmp_path, out)); + absl::Status renamed = env->RenameFile(tmp_path, std::string(path)); + if (!renamed.ok()) { + env->DeleteFile(tmp_path).IgnoreError(); + return renamed; + } + return absl::OkStatus(); } absl::StatusOr @@ -305,13 +354,22 @@ AOTInterceptionPjrtClient::PackArtifactForInnerClient() { ExecutableAndOptionsProto executable_and_options = std::move(*unpacked.mutable_executable_and_options()); - // Deriving the platform from the artifact (rather than the live inner client) - // to keep the un/packing pure. - const AOTTestPlatform platform = - unpacked.backend_executable_case() == - HumanReadableAotExecutable::kCpuExecutable - ? AOTTestPlatform::kCpu - : AOTTestPlatform::kGpu; + // Derived from the artifact, not the live client, to keep un/packing pure. An + // unset oneof means a malformed golden, not one of the platforms. + AOTTestPlatform platform; + switch (unpacked.backend_executable_case()) { + case HumanReadableAotExecutable::kCpuExecutable: + platform = AOTTestPlatform::kCpu; + break; + case HumanReadableAotExecutable::kGpuExecutable: + platform = AOTTestPlatform::kGpu; + break; + case HumanReadableAotExecutable::BACKEND_EXECUTABLE_NOT_SET: + return absl::InvalidArgumentError(absl::StrCat( + "AOTInterceptionPjrtClient: golden artifact has no backend " + "executable set: ", + artifact_path_)); + } if (platform == AOTTestPlatform::kGpu) { std::string serialized_gpu_exec; ABSL_RETURN_IF_ERROR(WriteSplitGpuExecutable( @@ -349,76 +407,66 @@ AOTInterceptionPjrtClient::PackArtifactForInnerClient() { absl::StatusOr> AOTInterceptionPjrtClient::Compile(const XlaComputation& computation, CompileOptions options) { - if (mode_ == AOTTestMode::kBackwardsCompatibility) { - VLOG(1) << "AOTInterceptionPjrtClient: Intercepting Compile in " - "kBackwardsCompatibility mode for computation [" - << computation.name() - << "]. Deserializing executable instead of recompiling."; - ABSL_ASSIGN_OR_RETURN(std::string serialized, PackArtifactForInnerClient()); - VLOG(1) << "AOTInterceptionPjrtClient: Calling " - "inner_client_->DeserializeExecutable."; - return inner_client_->DeserializeExecutable(serialized, std::move(options)); + switch (mode_) { + case AOTTestMode::kBackwardsCompatibility: { + VLOG(1) << "AOTInterceptionPjrtClient: Intercepting Compile in " + "kBackwardsCompatibility mode for computation [" + << computation.name() + << "]. Deserializing executable instead of recompiling."; + ABSL_ASSIGN_OR_RETURN(std::string serialized, PackArtifactForInnerClient()); + VLOG(1) << "AOTInterceptionPjrtClient: Calling " + "inner_client_->DeserializeExecutable."; + return inner_client_->DeserializeExecutable(serialized, + std::move(options)); + } + case AOTTestMode::kUpdateGolden: + case AOTTestMode::kGoldenVerification: { + ApplyAotDeterminismDebugOptions( + *options.executable_build_options.mutable_debug_options()); + ABSL_ASSIGN_OR_RETURN(std::unique_ptr exec, + inner_client_->Compile(computation, std::move(options))); + TF_RET_CHECK(exec != nullptr) << "Compile() returned nullptr"; + ABSL_RETURN_IF_ERROR(mode_ == AOTTestMode::kUpdateGolden + ? DumpGolden(*exec) + : VerifyAgainstGolden(*exec)); + return exec; + } } - if (mode_ == AOTTestMode::kGoldenVerification) { - VLOG(1) << "AOTInterceptionPjrtClient: Compile called in " - "kGoldenVerification mode for computation [" - << computation.name() - << "]. Compiling fresh and verifying against golden."; - options.executable_build_options.mutable_debug_options() - ->set_xla_gpu_exclude_nondeterministic_ops(true); - ABSL_ASSIGN_OR_RETURN(std::unique_ptr exec, - inner_client_->Compile(computation, std::move(options))); - - TF_RET_CHECK(exec != nullptr) << "Compile() returned nullptr"; - ABSL_RETURN_IF_ERROR(VerifyAgainstGolden(*exec)); - VLOG(1) << "AOTInterceptionPjrtClient: Golden Verification successful " - << "for [" << computation.name() << "]"; - - return exec; - } - - return absl::InternalError( - "Unknown AOTTestMode in AOTInterceptionPjrtClient::Compile"); + CHECK(false) << "Unsupported AOT test mode!"; } absl::StatusOr> AOTInterceptionPjrtClient::CompileAndLoad(const XlaComputation& computation, CompileOptions options) { - if (mode_ == AOTTestMode::kBackwardsCompatibility) { - VLOG(1) << "AOTInterceptionPjrtClient: Intercepting CompileAndLoad in " - "kBackwardsCompatibility mode for computation [" - << computation.name() - << "]. Deserializing and loading executable instead of " - "recompiling."; - ABSL_ASSIGN_OR_RETURN(std::string serialized, PackArtifactForInnerClient()); - VLOG(1) << "AOTInterceptionPjrtClient: Calling " - "inner_client_->LoadSerializedExecutable."; - return inner_client_->LoadSerializedExecutable( - serialized, std::move(options), LoadOptions()); + switch (mode_) { + case AOTTestMode::kBackwardsCompatibility: { + VLOG(1) << "AOTInterceptionPjrtClient: Intercepting CompileAndLoad in " + "kBackwardsCompatibility mode for computation [" + << computation.name() + << "]. Deserializing and loading executable instead of " + "recompiling."; + ABSL_ASSIGN_OR_RETURN(std::string serialized, PackArtifactForInnerClient()); + VLOG(1) << "AOTInterceptionPjrtClient: Calling " + "inner_client_->LoadSerializedExecutable."; + return inner_client_->LoadSerializedExecutable( + serialized, std::move(options), LoadOptions()); + } + case AOTTestMode::kUpdateGolden: + case AOTTestMode::kGoldenVerification: { + ApplyAotDeterminismDebugOptions( + *options.executable_build_options.mutable_debug_options()); + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr exec, + inner_client_->CompileAndLoad(computation, std::move(options))); + PjRtExecutable* fresh_exec = exec->GetExecutable(); + TF_RET_CHECK(fresh_exec != nullptr) << "GetExecutable() returned nullptr"; + ABSL_RETURN_IF_ERROR(mode_ == AOTTestMode::kUpdateGolden + ? DumpGolden(*fresh_exec) + : VerifyAgainstGolden(*fresh_exec)); + return exec; + } } - if (mode_ == AOTTestMode::kGoldenVerification) { - VLOG(1) << "AOTInterceptionPjrtClient: CompileAndLoad called in " - "kGoldenVerification mode for computation [" - << computation.name() - << "]. Compiling fresh and verifying against golden."; - options.executable_build_options.mutable_debug_options() - ->set_xla_gpu_exclude_nondeterministic_ops(true); - ABSL_ASSIGN_OR_RETURN( - std::unique_ptr exec, - inner_client_->CompileAndLoad(computation, std::move(options))); - - PjRtExecutable* fresh_exec = exec->GetExecutable(); - TF_RET_CHECK(fresh_exec != nullptr) << "GetExecutable() returned nullptr"; - - ABSL_RETURN_IF_ERROR(VerifyAgainstGolden(*fresh_exec)); - VLOG(1) << "AOTInterceptionPjrtClient: Golden Verification successful " - << "for [" << computation.name() << "]"; - - return exec; - } - - return absl::InternalError( - "Unknown AOTTestMode in AOTInterceptionPjrtClient::CompileAndLoad"); + CHECK(false) << "Unsupported AOT test mode!"; } absl::StatusOr AOTInterceptionPjrtClient::LookupDevice( diff --git a/third_party/xla/xla/tests/aot_interception_pjrt_client.h b/third_party/xla/xla/tests/aot_interception_pjrt_client.h index 0e9ff7a378a649..b6f2ac2b66d36f 100644 --- a/third_party/xla/xla/tests/aot_interception_pjrt_client.h +++ b/third_party/xla/xla/tests/aot_interception_pjrt_client.h @@ -35,6 +35,7 @@ limitations under the License. #include "xla/runtime/device_id.h" #include "xla/service/device_assignment.h" #include "xla/util/split_proto/human_readable_aot_executable.pb.h" +#include "xla/xla.pb.h" #include "xla/xla_data.pb.h" namespace xla { @@ -43,6 +44,7 @@ namespace xla { enum class AOTTestMode { kGoldenVerification, kBackwardsCompatibility, + kUpdateGolden, }; enum class AOTTestPlatform { @@ -82,6 +84,14 @@ class AOTInterceptionPjrtClient : public PjRtClient { // Unpacks a serialized PjRtExecutable into the human-readable proto form. static absl::StatusOr DeserializeToHumanReadable( absl::string_view serialized, AOTTestPlatform platform); + + static absl::Status WriteGoldenTextProto( + const HumanReadableAotExecutable& unpacked, absl::string_view path); + + // Modifies `debug_options` to configure deterministic AOT compilation for + // golden generation and verification. + static void ApplyAotDeterminismDebugOptions(DebugOptions& debug_options); + absl::StatusOr> Compile( const XlaComputation& computation, CompileOptions options) override; @@ -161,9 +171,9 @@ class AOTInterceptionPjrtClient : public PjRtClient { private: absl::Status VerifyAgainstGolden(const PjRtExecutable& fresh_executable); + absl::Status DumpGolden(const PjRtExecutable& fresh_executable); - // Loads the artifact at artifact_path_ and parses it into its human-readable - // proto form. Used internally by the compile and verification paths. + // Loads the artifact at artifact_path_ into its human-readable proto form. absl::StatusOr LoadHumanReadableArtifact(); std::unique_ptr inner_client_; diff --git a/third_party/xla/xla/tools/hlo_decomposer.cc b/third_party/xla/xla/tools/hlo_decomposer.cc index 0172b1c8ff1b23..1412140480af3e 100644 --- a/third_party/xla/xla/tools/hlo_decomposer.cc +++ b/third_party/xla/xla/tools/hlo_decomposer.cc @@ -171,7 +171,15 @@ std::unique_ptr ExtractCollectiveOperationsIntoNewModule( HloInstruction::CreateTuple(result_instructions); builder.AddInstruction(std::move(tuple_instruction)); } - new_hlo_module->AddEntryComputationWithLayouts(builder.Build()); + HloComputation* entry_computation = + new_hlo_module->AddEntryComputationWithLayouts(builder.Build()); + // Preserve the execution thread so the extracted module stays internally + // consistent: a cloned instruction embedded below the entry (e.g. a fusion + // wrapping an op that lives on a non-"main" thread) still carries its + // original execution_thread, and HloVerifier requires the calling + // computation's thread to match. + entry_computation->SetExecutionThread( + first_instruction.parent()->execution_thread()); return new_hlo_module; } @@ -194,7 +202,10 @@ std::unique_ptr ExtractInstructionIntoNewModule( std::unique_ptr new_instruction = hlo.CloneWithNewOperands(hlo.shape(), new_operands, &clone_context); builder.AddInstruction(std::move(new_instruction)); - new_hlo_module->AddEntryComputationWithLayouts(builder.Build()); + HloComputation* entry_computation = + new_hlo_module->AddEntryComputationWithLayouts(builder.Build()); + // See the comment in ExtractCollectiveOperationsIntoNewModule above. + entry_computation->SetExecutionThread(hlo.parent()->execution_thread()); return new_hlo_module; } @@ -247,7 +258,10 @@ std::unique_ptr ExtractProducerConsumersIntoNewModule( builder.AddInstruction(HloInstruction::CreateTuple(consumer_outputs)); } - new_hlo_module->AddEntryComputationWithLayouts(builder.Build()); + HloComputation* entry_computation = + new_hlo_module->AddEntryComputationWithLayouts(builder.Build()); + // See the comment in ExtractCollectiveOperationsIntoNewModule above. + entry_computation->SetExecutionThread(producer.parent()->execution_thread()); return new_hlo_module; } diff --git a/third_party/xla/xla/tools/hlo_decomposer_test.cc b/third_party/xla/xla/tools/hlo_decomposer_test.cc index 339aabd91f8cfc..1ed819c11bf2ff 100644 --- a/third_party/xla/xla/tools/hlo_decomposer_test.cc +++ b/third_party/xla/xla/tools/hlo_decomposer_test.cc @@ -204,6 +204,75 @@ CHECK-THEN: ROOT %tuple.1 = (f32[10,10]{1,0}, f32[10,10]{1,0}) tuple(%add.1, %su EXPECT_TRUE(check_result); } +TEST_F(HloDecomposerTest, + ExtractInstructionIntoNewModulePreservesExecutionThread) { + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( +HloModule module + +fused_computation { + p0 = f32[8] parameter(0) + ROOT neg = f32[8] negate(p0) +}, execution_thread="parallel" + +parallel_comp { + p0 = f32[8] parameter(0) + ROOT loop_fusion = f32[8] fusion(p0), kind=kLoop, calls=fused_computation +}, execution_thread="parallel" + +ENTRY main { + p0 = f32[8] parameter(0) + call-start = ((f32[8]), f32[8]) call-start(p0), async_execution_thread="parallel", to_apply=parallel_comp + ROOT call-done = f32[8] call-done(call-start) +})")); + HloComputation* parallel_comp = + module->GetComputationWithName("parallel_comp"); + ASSERT_NE(parallel_comp, nullptr); + HloInstruction* loop_fusion = + parallel_comp->GetInstructionWithName("loop_fusion"); + ASSERT_NE(loop_fusion, nullptr); + + std::unique_ptr new_module = + ExtractInstructionIntoNewModule(*loop_fusion); + + EXPECT_EQ(new_module->entry_computation()->execution_thread(), "parallel"); +} + +TEST_F(HloDecomposerTest, + ExtractProducerConsumersIntoNewModulePreservesExecutionThread) { + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( +HloModule module + +fused_computation { + p0 = f32[8] parameter(0) + ROOT neg = f32[8] negate(p0) +}, execution_thread="parallel" + +parallel_comp { + p0 = f32[8] parameter(0) + loop_fusion = f32[8] fusion(p0), kind=kLoop, calls=fused_computation + ROOT result = f32[8] add(loop_fusion, p0) +}, execution_thread="parallel" + +ENTRY main { + p0 = f32[8] parameter(0) + call-start = ((f32[8]), f32[8]) call-start(p0), async_execution_thread="parallel", to_apply=parallel_comp + ROOT call-done = f32[8] call-done(call-start) +})")); + HloComputation* parallel_comp = + module->GetComputationWithName("parallel_comp"); + ASSERT_NE(parallel_comp, nullptr); + HloInstruction* loop_fusion = + parallel_comp->GetInstructionWithName("loop_fusion"); + ASSERT_NE(loop_fusion, nullptr); + + std::unique_ptr new_module = + ExtractProducerConsumersIntoNewModule(*loop_fusion); + + EXPECT_EQ(new_module->entry_computation()->execution_thread(), "parallel"); +} + TEST_F(HloDecomposerTest, ExtractProducerConsumersIntoNewModuleSingleConsumer) { TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(R"( diff --git a/third_party/xla/xla/xla.default.bzl b/third_party/xla/xla/xla.default.bzl index d27ded2e99733a..e134a1b7f32e83 100644 --- a/third_party/xla/xla/xla.default.bzl +++ b/third_party/xla/xla/xla.default.bzl @@ -67,7 +67,6 @@ _XLA_SHARED_OBJECT_SENSITIVE_DEPS = [ Label("//xla/service:metrics_proto_cc_impl"), Label("//xla/service/gpu:backend_configs_cc_impl"), Label("//xla/service/gpu/model:hlo_op_profile_proto_cc_impl"), - Label("//xla/service/memory_space_assignment:memory_space_assignment_proto_cc_impl"), Label("//xla/stream_executor:device_description_proto_cc_impl"), Label("//xla/stream_executor:stream_executor_impl"), Label("//xla/stream_executor/cuda:cuda_compute_capability_proto_cc_impl"),