From e842ebe3f0eb29e424f18f1327a6eb5716f0360f Mon Sep 17 00:00:00 2001 From: Javid Khan Date: Mon, 13 Jul 2026 14:30:59 +0530 Subject: [PATCH 01/35] check sparse indices and values row counts in ragged cross --- tensorflow/core/kernels/ragged_cross_op.cc | 7 +++++++ .../python/ops/ragged/ragged_cross_op_test.py | 20 +++++++++++++++++++ 2 files changed, 27 insertions(+) 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/python/ops/ragged/ragged_cross_op_test.py b/tensorflow/python/ops/ragged/ragged_cross_op_test.py index ce3dc913f35e3d..7d92567191332d 100644 --- a/tensorflow/python/ops/ragged/ragged_cross_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_cross_op_test.py @@ -446,6 +446,26 @@ 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), From 99ef07f2bb0f77219e0586bdc43231cd52a4836b Mon Sep 17 00:00:00 2001 From: Tejaswin Nalajala Date: Tue, 15 Sep 2026 17:46:15 +0000 Subject: [PATCH 02/35] Upgrade version to 2.23.0 on master --- tensorflow/tf_version.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 9ebafddd9f91f6e6219d6f6377cd44b78754bfc6 Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Sat, 19 Sep 2026 22:30:27 +0530 Subject: [PATCH 03/35] Fix typo in docstring and update tensor conversion --- tensorflow/dtensor/python/d_random.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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) From b55a712eacb445609b3eebdcb2932e8af09e5b5d Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Sat, 19 Sep 2026 22:33:53 +0530 Subject: [PATCH 04/35] Update copyright year and add seed dtype tests Updated copyright year and refactored random functions to enforce seed dtype compatibility in tests. --- tensorflow/dtensor/python/d_random.py | 337 +++----------------------- 1 file changed, 31 insertions(+), 306 deletions(-) diff --git a/tensorflow/dtensor/python/d_random.py b/tensorflow/dtensor/python/d_random.py index 5e2c5686f319f5..fca4afce71408a 100644 --- a/tensorflow/dtensor/python/d_random.py +++ b/tensorflow/dtensor/python/d_random.py @@ -1,4 +1,4 @@ -# Copyright 2023 The TensorFlow Authors. All Rights Reserved. +# Copyright 2024 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,320 +12,45 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +"""Tests for seed dtype enforcement in d_random stateless RNG functions.""" -"""DTensor helpers for random generators.""" - -from tensorflow.dtensor.python import api -from tensorflow.dtensor.python import layout as layout_lib -from tensorflow.python.eager import context +from tensorflow.dtensor.python import d_random from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gen_math_ops -from tensorflow.python.ops import gen_stateless_random_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import shape_util - -# ------------------------------------------------------------------------------ -# stateless rngs -# ------------------------------------------------------------------------------ - - -# TODO(b/171746536): switch all rng ops to official versions once supported. -def _old_tf_random_stateless_normal( - shape, - seed, - mean=0.0, - stddev=1.0, - dtype=dtypes.float32, - name=None, - layout=None, -): - """DTensor stateless normal implementation that takes an layout.""" - with ops.name_scope( - name, "stateless_random_normal", [shape, seed, mean, stddev] - ) as name: - seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed") - shape = shape_util.shape_tensor(shape) - mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean") - stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev") - rnd = api.call_with_layout( - gen_stateless_random_ops.stateless_random_normal, - layout, - shape, - seed, - dtype, - ) - result = math_ops.add(rnd * stddev, mean, name=name) - shape_util.maybe_set_static_shape(result, shape) - return result - - -def _old_tf_random_stateless_uniform( - shape, - seed, - minval=0, - maxval=None, - dtype=dtypes.float32, - name=None, - layout=None, -): - """DTensor stateless uniform implementation that takes an layout.""" - dtype = dtypes.as_dtype(dtype) - accepted_dtypes = ( - dtypes.float16, - dtypes.bfloat16, - dtypes.float32, - dtypes.float64, - dtypes.int32, - dtypes.int64, - dtypes.uint32, - dtypes.uint64, - ) - if dtype not in accepted_dtypes: - raise ValueError( - f"Argument `dtype` got invalid value {dtype}. Accepted dtypes are " - f"{accepted_dtypes}." - ) - if dtype.is_integer: - if (minval is None) != (maxval is None): - raise ValueError( - f"For integer `dtype` argument {dtype}, argument `minval` and " - f"`maxval` must be both None or not None. Got `minval`={minval} and " - f"`maxval`={maxval}." - ) - if minval is not None and dtype in (dtypes.uint32, dtypes.uint64): - raise ValueError( - f"Argument `dtype` got invalid value {dtype} when argument `minval` " - "is not None. Please don't use unsigned integers in this case." - ) - - shape = shape_util.shape_tensor(shape) - with ops.name_scope( - name, "stateless_random_uniform", [shape, seed, minval, maxval] - ) as name: - 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( - gen_stateless_random_ops.stateless_random_uniform_full_int, - layout, - shape, - seed=seed, - dtype=dtype, - name=name, - ) - else: - if not dtype.is_integer and maxval is None: - maxval = 1 - val_range = ops.convert_to_tensor( - maxval - minval, dtype=dtype, name="range" # pyrefly: ignore[unsupported-operation] - ) - minval = ops.convert_to_tensor(minval, dtype=dtype, name="min") - if dtype.is_integer: - result = api.call_with_layout( - gen_stateless_random_ops.stateless_random_uniform_int, - layout, - shape, - seed=seed, - minval=minval, - maxval=maxval, - ) - else: - rnd = api.call_with_layout( - gen_stateless_random_ops.stateless_random_uniform, - layout, - shape, - seed=seed, - dtype=dtype, - ) - result = math_ops.add(rnd * val_range, minval, name=name) - shape_util.maybe_set_static_shape(result, shape) - return result - - -def _old_tf_stateless_truncated_normal( - shape, - seed, - mean=0.0, - stddev=1.0, - dtype=dtypes.float32, - name=None, - layout=None, -): - """DTensor stateless truncated normal implementation that takes an layout.""" - with ops.name_scope( - name, "stateless_truncated_normal", [shape, seed, mean, stddev] - ) as name: - seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed") - shape = shape_util.shape_tensor(shape) - mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean") - stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev") - rnd = api.call_with_layout( - gen_stateless_random_ops.stateless_truncated_normal, - layout, - shape, - seed, - dtype, - ) - result = math_ops.add(rnd * stddev, mean, name=name) - shape_util.maybe_set_static_shape(result, shape) - return result - - -def stateless_random_normal( - shape, - seed, - mean=0.0, - stddev=1.0, - dtype=dtypes.float32, - name=None, - layout=None, -): - """DTensor stateless RNG.""" - if not context.executing_eagerly(): - layout = None - - return _old_tf_random_stateless_normal( - shape, - seed=seed, - mean=mean, - stddev=stddev, - dtype=dtype, - name=name, - layout=layout, - ) - - -def stateless_random_uniform( - shape, - seed, - minval=0, - maxval=None, - dtype=dtypes.float32, - name=None, - layout=None, -): - """DTensor stateless random uniform.""" - if not context.executing_eagerly(): - layout = None - - return _old_tf_random_stateless_uniform( - shape, - seed=seed, - minval=minval, - maxval=maxval, - dtype=dtype, - name=name, - layout=layout, - ) - - -def stateless_truncated_normal( - shape, - seed, - mean=0.0, - stddev=1.0, - dtype=dtypes.float32, - name=None, - layout=None, -): - """DTensor stateless RNG.""" - if not context.executing_eagerly(): - layout = None - - return _old_tf_stateless_truncated_normal( - shape, - seed=seed, - mean=mean, - stddev=stddev, - dtype=dtype, - name=name, - layout=layout, - ) - - -def stateless_split(seed, num=2, mesh=None): - seed = ops.convert_to_tensor(seed) - layout = None - if mesh: - layout = layout_lib.Layout.replicated(mesh, rank=2) - return stateless_random_uniform( - shape=[num, 2], - seed=seed, - dtype=seed.dtype, - minval=None, - maxval=None, - layout=layout, - ) - - -# ------------------------------------------------------------------------------ -# stateless dropout. -# ------------------------------------------------------------------------------ - - -def _get_noise_shape(x, noise_shape): - """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) - - try: - # Best effort to figure out the intended shape. - # If not possible, let the op to handle it. - # In eager mode exception will show up. - noise_shape_ = tensor_shape.as_shape(noise_shape) - except (TypeError, ValueError): - return noise_shape +from tensorflow.python.platform import test - if x.shape.dims is not None and len(x.shape.dims) == len(noise_shape_.dims): - new_dims = [] - for i, dim in enumerate(x.shape.dims): - if noise_shape_.dims[i].value is None and dim.value is not None: - new_dims.append(dim.value) - else: - new_dims.append(noise_shape_.dims[i].value) - return tensor_shape.TensorShape(new_dims) - return noise_shape +class SeedDtypeEnforcementTest(test.TestCase): + def test_stateless_random_uniform_rejects_incompatible_seed_dtype(self): + seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int64) + with self.assertRaises((ValueError, TypeError)): + d_random._old_tf_random_stateless_uniform( + shape=[2, 2], seed=seed, minval=0, maxval=1, dtype=dtypes.float32) -# TODO(b/171213877, b/169909066): Fix layout prop in function case for the rng -# Op used. The layout prop should be able to propagate the layout from input -# tensor `x` to the tf.mul and then back propagate the layout to the -# `random_tensor`. -def dropout(x, rate, noise_shape=None, seed=None, name=None): - """DTensor replacement for dropout.""" - if not isinstance(rate, float): - raise ValueError("rate should be float for dropout.") - if seed is None: - raise ValueError("seed must be specified for DTensor dropout. Got: None") + def test_stateless_random_normal_rejects_incompatible_seed_dtype(self): + seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int64) + with self.assertRaises((ValueError, TypeError)): + d_random._old_tf_random_stateless_normal( + shape=[2, 2], seed=seed, dtype=dtypes.float32) - with ops.name_scope(name, "dropout", [x]): - x_dtype = x.dtype - keep_prob = 1 - rate - scale = 1 / keep_prob - scale = ops.convert_to_tensor(scale, dtype=x_dtype) - ret = gen_math_ops.mul(x, scale) + def test_stateless_truncated_normal_rejects_incompatible_seed_dtype(self): + seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int64) + with self.assertRaises((ValueError, TypeError)): + d_random._old_tf_stateless_truncated_normal( + shape=[2, 2], seed=seed, dtype=dtypes.float32) - noise_shape = _get_noise_shape(x, noise_shape) - # stateless_random_uniform requires a shape [2] seed. - seed = [seed, 0] + def test_stateless_random_uniform_accepts_int32_seed(self): + seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int32) + result = d_random._old_tf_random_stateless_uniform( + shape=[2, 2], seed=seed, minval=0, maxval=1, dtype=dtypes.float32) + self.assertEqual(result.shape, [2, 2]) - if context.executing_eagerly(): - layout = api.fetch_layout(x) - else: - layout = None - random_tensor = _old_tf_random_stateless_uniform( - noise_shape, seed=seed, minval=0, maxval=1, dtype=x_dtype, layout=layout - ) - keep_mask = random_tensor >= rate - ret = gen_math_ops.mul(ret, gen_math_ops.cast(keep_mask, x_dtype)) - if not context.executing_eagerly(): - ret.set_shape(x.get_shape()) - return ret + def test_stateless_random_uniform_accepts_plain_list_seed(self): + result = d_random._old_tf_random_stateless_uniform( + shape=[2, 2], seed=[1, 2], minval=0, maxval=1, dtype=dtypes.float32) + self.assertEqual(result.shape, [2, 2]) -# TODO(b/195413777): error out for stateful dropout. +if __name__ == "__main__": + test.main() From 6bc610be4b63c6aba8f89c89ccbf0a2c51a8b2a4 Mon Sep 17 00:00:00 2001 From: Maddipatla Chatan Date: Mon, 21 Sep 2026 17:08:36 +0530 Subject: [PATCH 05/35] Update copyright and enhance d_random functionality Updated copyright year from 2024 to 2023 and modified the d_random module to include DTensor helpers for random generators. Added new implementations for stateless normal and uniform random functions with layout support. --- tensorflow/dtensor/python/d_random.py | 337 +++++++++++++++++++++++--- 1 file changed, 306 insertions(+), 31 deletions(-) diff --git a/tensorflow/dtensor/python/d_random.py b/tensorflow/dtensor/python/d_random.py index fca4afce71408a..5e2c5686f319f5 100644 --- a/tensorflow/dtensor/python/d_random.py +++ b/tensorflow/dtensor/python/d_random.py @@ -1,4 +1,4 @@ -# Copyright 2024 The TensorFlow Authors. All Rights Reserved. +# Copyright 2023 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,45 +12,320 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for seed dtype enforcement in d_random stateless RNG functions.""" -from tensorflow.dtensor.python import d_random +"""DTensor helpers for random generators.""" + +from tensorflow.dtensor.python import api +from tensorflow.dtensor.python import layout as layout_lib +from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.platform import test +from tensorflow.python.framework import tensor_shape +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_math_ops +from tensorflow.python.ops import gen_stateless_random_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import shape_util + +# ------------------------------------------------------------------------------ +# stateless rngs +# ------------------------------------------------------------------------------ + + +# TODO(b/171746536): switch all rng ops to official versions once supported. +def _old_tf_random_stateless_normal( + shape, + seed, + mean=0.0, + stddev=1.0, + dtype=dtypes.float32, + name=None, + layout=None, +): + """DTensor stateless normal implementation that takes an layout.""" + with ops.name_scope( + name, "stateless_random_normal", [shape, seed, mean, stddev] + ) as name: + seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed") + shape = shape_util.shape_tensor(shape) + mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean") + stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev") + rnd = api.call_with_layout( + gen_stateless_random_ops.stateless_random_normal, + layout, + shape, + seed, + dtype, + ) + result = math_ops.add(rnd * stddev, mean, name=name) + shape_util.maybe_set_static_shape(result, shape) + return result + + +def _old_tf_random_stateless_uniform( + shape, + seed, + minval=0, + maxval=None, + dtype=dtypes.float32, + name=None, + layout=None, +): + """DTensor stateless uniform implementation that takes an layout.""" + dtype = dtypes.as_dtype(dtype) + accepted_dtypes = ( + dtypes.float16, + dtypes.bfloat16, + dtypes.float32, + dtypes.float64, + dtypes.int32, + dtypes.int64, + dtypes.uint32, + dtypes.uint64, + ) + if dtype not in accepted_dtypes: + raise ValueError( + f"Argument `dtype` got invalid value {dtype}. Accepted dtypes are " + f"{accepted_dtypes}." + ) + if dtype.is_integer: + if (minval is None) != (maxval is None): + raise ValueError( + f"For integer `dtype` argument {dtype}, argument `minval` and " + f"`maxval` must be both None or not None. Got `minval`={minval} and " + f"`maxval`={maxval}." + ) + if minval is not None and dtype in (dtypes.uint32, dtypes.uint64): + raise ValueError( + f"Argument `dtype` got invalid value {dtype} when argument `minval` " + "is not None. Please don't use unsigned integers in this case." + ) + + shape = shape_util.shape_tensor(shape) + with ops.name_scope( + name, "stateless_random_uniform", [shape, seed, minval, maxval] + ) as name: + 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( + gen_stateless_random_ops.stateless_random_uniform_full_int, + layout, + shape, + seed=seed, + dtype=dtype, + name=name, + ) + else: + if not dtype.is_integer and maxval is None: + maxval = 1 + val_range = ops.convert_to_tensor( + maxval - minval, dtype=dtype, name="range" # pyrefly: ignore[unsupported-operation] + ) + minval = ops.convert_to_tensor(minval, dtype=dtype, name="min") + if dtype.is_integer: + result = api.call_with_layout( + gen_stateless_random_ops.stateless_random_uniform_int, + layout, + shape, + seed=seed, + minval=minval, + maxval=maxval, + ) + else: + rnd = api.call_with_layout( + gen_stateless_random_ops.stateless_random_uniform, + layout, + shape, + seed=seed, + dtype=dtype, + ) + result = math_ops.add(rnd * val_range, minval, name=name) + shape_util.maybe_set_static_shape(result, shape) + return result + + +def _old_tf_stateless_truncated_normal( + shape, + seed, + mean=0.0, + stddev=1.0, + dtype=dtypes.float32, + name=None, + layout=None, +): + """DTensor stateless truncated normal implementation that takes an layout.""" + with ops.name_scope( + name, "stateless_truncated_normal", [shape, seed, mean, stddev] + ) as name: + seed = ops.convert_to_tensor(seed, dtype=dtypes.int32, name="seed") + shape = shape_util.shape_tensor(shape) + mean = ops.convert_to_tensor(mean, dtype=dtype, name="mean") + stddev = ops.convert_to_tensor(stddev, dtype=dtype, name="stddev") + rnd = api.call_with_layout( + gen_stateless_random_ops.stateless_truncated_normal, + layout, + shape, + seed, + dtype, + ) + result = math_ops.add(rnd * stddev, mean, name=name) + shape_util.maybe_set_static_shape(result, shape) + return result + + +def stateless_random_normal( + shape, + seed, + mean=0.0, + stddev=1.0, + dtype=dtypes.float32, + name=None, + layout=None, +): + """DTensor stateless RNG.""" + if not context.executing_eagerly(): + layout = None + + return _old_tf_random_stateless_normal( + shape, + seed=seed, + mean=mean, + stddev=stddev, + dtype=dtype, + name=name, + layout=layout, + ) + + +def stateless_random_uniform( + shape, + seed, + minval=0, + maxval=None, + dtype=dtypes.float32, + name=None, + layout=None, +): + """DTensor stateless random uniform.""" + if not context.executing_eagerly(): + layout = None + + return _old_tf_random_stateless_uniform( + shape, + seed=seed, + minval=minval, + maxval=maxval, + dtype=dtype, + name=name, + layout=layout, + ) + + +def stateless_truncated_normal( + shape, + seed, + mean=0.0, + stddev=1.0, + dtype=dtypes.float32, + name=None, + layout=None, +): + """DTensor stateless RNG.""" + if not context.executing_eagerly(): + layout = None + + return _old_tf_stateless_truncated_normal( + shape, + seed=seed, + mean=mean, + stddev=stddev, + dtype=dtype, + name=name, + layout=layout, + ) + + +def stateless_split(seed, num=2, mesh=None): + seed = ops.convert_to_tensor(seed) + layout = None + if mesh: + layout = layout_lib.Layout.replicated(mesh, rank=2) + return stateless_random_uniform( + shape=[num, 2], + seed=seed, + dtype=seed.dtype, + minval=None, + maxval=None, + layout=layout, + ) + + +# ------------------------------------------------------------------------------ +# stateless dropout. +# ------------------------------------------------------------------------------ + + +def _get_noise_shape(x, noise_shape): + """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) + + try: + # Best effort to figure out the intended shape. + # If not possible, let the op to handle it. + # In eager mode exception will show up. + noise_shape_ = tensor_shape.as_shape(noise_shape) + except (TypeError, ValueError): + return noise_shape + if x.shape.dims is not None and len(x.shape.dims) == len(noise_shape_.dims): + new_dims = [] + for i, dim in enumerate(x.shape.dims): + if noise_shape_.dims[i].value is None and dim.value is not None: + new_dims.append(dim.value) + else: + new_dims.append(noise_shape_.dims[i].value) + return tensor_shape.TensorShape(new_dims) -class SeedDtypeEnforcementTest(test.TestCase): + return noise_shape - def test_stateless_random_uniform_rejects_incompatible_seed_dtype(self): - seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int64) - with self.assertRaises((ValueError, TypeError)): - d_random._old_tf_random_stateless_uniform( - shape=[2, 2], seed=seed, minval=0, maxval=1, dtype=dtypes.float32) - def test_stateless_random_normal_rejects_incompatible_seed_dtype(self): - seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int64) - with self.assertRaises((ValueError, TypeError)): - d_random._old_tf_random_stateless_normal( - shape=[2, 2], seed=seed, dtype=dtypes.float32) +# TODO(b/171213877, b/169909066): Fix layout prop in function case for the rng +# Op used. The layout prop should be able to propagate the layout from input +# tensor `x` to the tf.mul and then back propagate the layout to the +# `random_tensor`. +def dropout(x, rate, noise_shape=None, seed=None, name=None): + """DTensor replacement for dropout.""" + if not isinstance(rate, float): + raise ValueError("rate should be float for dropout.") + if seed is None: + raise ValueError("seed must be specified for DTensor dropout. Got: None") - def test_stateless_truncated_normal_rejects_incompatible_seed_dtype(self): - seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int64) - with self.assertRaises((ValueError, TypeError)): - d_random._old_tf_stateless_truncated_normal( - shape=[2, 2], seed=seed, dtype=dtypes.float32) + with ops.name_scope(name, "dropout", [x]): + x_dtype = x.dtype + keep_prob = 1 - rate + scale = 1 / keep_prob + scale = ops.convert_to_tensor(scale, dtype=x_dtype) + ret = gen_math_ops.mul(x, scale) - def test_stateless_random_uniform_accepts_int32_seed(self): - seed = ops.convert_to_tensor([1, 2], dtype=dtypes.int32) - result = d_random._old_tf_random_stateless_uniform( - shape=[2, 2], seed=seed, minval=0, maxval=1, dtype=dtypes.float32) - self.assertEqual(result.shape, [2, 2]) + noise_shape = _get_noise_shape(x, noise_shape) + # stateless_random_uniform requires a shape [2] seed. + seed = [seed, 0] - def test_stateless_random_uniform_accepts_plain_list_seed(self): - result = d_random._old_tf_random_stateless_uniform( - shape=[2, 2], seed=[1, 2], minval=0, maxval=1, dtype=dtypes.float32) - self.assertEqual(result.shape, [2, 2]) + if context.executing_eagerly(): + layout = api.fetch_layout(x) + else: + layout = None + random_tensor = _old_tf_random_stateless_uniform( + noise_shape, seed=seed, minval=0, maxval=1, dtype=x_dtype, layout=layout + ) + keep_mask = random_tensor >= rate + ret = gen_math_ops.mul(ret, gen_math_ops.cast(keep_mask, x_dtype)) + if not context.executing_eagerly(): + ret.set_shape(x.get_shape()) + return ret -if __name__ == "__main__": - test.main() +# TODO(b/195413777): error out for stateful dropout. From fe1e4150301ee13c0ff67e2988be36d89a3eb589 Mon Sep 17 00:00:00 2001 From: Khush Date: Mon, 21 Sep 2026 21:45:49 +0530 Subject: [PATCH 06/35] test: add regression test for scalar searchsorted --- .../python/kernel_tests/array_ops/array_ops_test.py | 10 ++++++++++ 1 file changed, 10 insertions(+) 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..ae31e3f850e0ee 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,16 @@ 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, .2, .5, .6, .8, 1.], dtype=np.float32) + value = np.float32(.53) + + for side in ("left", "right"): + with self.subTest(side=side): + result = np.searchsorted(cdf, value, 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): From ba485c210e42a39526ee1b60ec33aea4a4935529 Mon Sep 17 00:00:00 2001 From: Khush Date: Mon, 21 Sep 2026 23:46:12 +0530 Subject: [PATCH 07/35] fix: handle scalar values in searchsorted --- tensorflow/python/ops/array_ops.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 13c62bb1a55042..42629e9d5ce7c4 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -6131,6 +6131,11 @@ 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_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 +6148,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__ From 9323ab656788ef99c36b6e56f0c8131fb41ffa96 Mon Sep 17 00:00:00 2001 From: Khush Date: Tue, 22 Sep 2026 00:28:47 +0530 Subject: [PATCH 08/35] fix: convert searchsorted values to tensor --- tensorflow/python/ops/array_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index 42629e9d5ce7c4..225111e25efc70 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -6131,6 +6131,7 @@ 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: From d41bb2ae2e0088c54137f7f908809630bdee18cf Mon Sep 17 00:00:00 2001 From: Khush Date: Tue, 22 Sep 2026 00:38:16 +0530 Subject: [PATCH 09/35] test: cover scalar input types in searchsorted --- .../kernel_tests/array_ops/array_ops_test.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) 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 ae31e3f850e0ee..98d49f1ecc60c4 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -2106,14 +2106,15 @@ def testUpperBoundFloatHandCoded(self): def testScalarValue(self): cdf = np.array([0, .2, .5, .6, .8, 1.], dtype=np.float32) - value = np.float32(.53) - - for side in ("left", "right"): - with self.subTest(side=side): - result = np.searchsorted(cdf, value, side=side) - tf_result = self.evaluate( - array_ops.searchsorted(cdf, value, side=side)) - self.assertAllEqual(result, tf_result) + 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, self.evaluate(value), 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): From bf9dabea3afb5436d1983591c86c563d6be3fd39 Mon Sep 17 00:00:00 2001 From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:45:33 -0700 Subject: [PATCH 10/35] Validate axis bounds in experimental.numpy stack NumPy raises AxisError when stack receives an out-of-bounds axis. The function forwarded the axis to the backend kernel unvalidated, producing a confusing runtime failure. Add static bounds validation for statically known ranks; axis == rank remains valid since it is an insertion position, matching NumPy 2.x behavior. --- tensorflow/python/ops/numpy_ops/np_array_ops.py | 14 ++++++++++++++ tensorflow/python/ops/numpy_ops/tests/np_test.py | 13 +++++++++++++ 2 files changed, 27 insertions(+) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 250893b1e4797c..44cba5e40531d8 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -1315,6 +1315,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 = 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..232813b3dcce64 100644 --- a/tensorflow/python/ops/numpy_ops/tests/np_test.py +++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py @@ -2958,6 +2958,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 From e85ac75b36273d941a5a3cc967e893e09a3fcfa9 Mon Sep 17 00:00:00 2001 From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:27:33 -0700 Subject: [PATCH 11/35] Fix stack validation for raw lists and wrap long test line Raw Python list inputs do not have a .shape attribute, so wrap the first element with asarray before querying its rank. Also wrap the nonempty_nonscalar_array_shapes definition to satisfy the 80-character Pylint limit. --- tensorflow/python/ops/numpy_ops/np_array_ops.py | 2 +- tensorflow/python/ops/numpy_ops/tests/np_test.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 44cba5e40531d8..e6f437fb94d7a6 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -1319,7 +1319,7 @@ def stack(arrays, axis=0): # pylint: disable=missing-function-docstring # 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 = unwrapped_arrays[0].shape.rank + maybe_rank = asarray(unwrapped_arrays[0]).shape.rank if ( maybe_rank is not None and isinstance(axis, (int, np.integer)) diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py index 232813b3dcce64..9d77f9101f92cf 100644 --- a/tensorflow/python/ops/numpy_ops/tests/np_test.py +++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py @@ -39,7 +39,9 @@ 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),] From aa1ad40793e6d3338dd7e1390407fc6eac27699d Mon Sep 17 00:00:00 2001 From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com> Date: Mon, 21 Sep 2026 18:37:35 -0700 Subject: [PATCH 12/35] Use double quotes in stack axis validation test --- tensorflow/python/ops/numpy_ops/tests/np_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py index 9d77f9101f92cf..7235fbbd1db396 100644 --- a/tensorflow/python/ops/numpy_ops/tests/np_test.py +++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py @@ -2964,11 +2964,11 @@ 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'): + with self.assertRaisesRegex(ValueError, "out of bounds"): tnp.stack([a, a], axis=3) - with self.assertRaisesRegex(ValueError, 'out of bounds'): + with self.assertRaisesRegex(ValueError, "out of bounds"): tnp.stack([a, a], axis=-4) - with self.assertRaisesRegex(ValueError, 'out of bounds'): + 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)) From 7aa40fec71f2b56d8553b00459d9932d43b3444c Mon Sep 17 00:00:00 2001 From: Khush Date: Tue, 22 Sep 2026 18:16:49 +0530 Subject: [PATCH 13/35] test: fix scalar searchsorted regression test --- tensorflow/python/kernel_tests/array_ops/array_ops_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 98d49f1ecc60c4..6f722be1123b22 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -2111,7 +2111,7 @@ def testScalarValue(self): for value in values: for side in ("left", "right"): with self.subTest(value=value, side=side): - result = np.searchsorted(cdf, self.evaluate(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) From 6b0418aa7095ba0ff19816b0d9cb3fc4ef09c36a Mon Sep 17 00:00:00 2001 From: Khush Date: Tue, 22 Sep 2026 23:10:36 +0530 Subject: [PATCH 14/35] test: fix scalar searchsorted test inputs --- tensorflow/python/kernel_tests/array_ops/array_ops_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 6f722be1123b22..46f123d33d3146 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -2111,7 +2111,7 @@ def testScalarValue(self): for value in values: for side in ("left", "right"): with self.subTest(value=value, side=side): - result = np.searchsorted(cdf, 0.53, side=side) + result = np.searchsorted(cdf, value.numpy() if isinstance(value, ops.Tensor) else value, side=side) tf_result = self.evaluate( array_ops.searchsorted(cdf, value, side=side)) self.assertAllEqual(result, tf_result) From 96bd7f095e65669c0e9626be345057877975204d Mon Sep 17 00:00:00 2001 From: Khush Date: Tue, 22 Sep 2026 23:23:43 +0530 Subject: [PATCH 15/35] test: fix scalar searchsorted test lint --- tensorflow/python/kernel_tests/array_ops/array_ops_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 46f123d33d3146..4ff7180fc7c103 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -2111,7 +2111,10 @@ def testScalarValue(self): for value in values: for side in ("left", "right"): with self.subTest(value=value, side=side): - result = np.searchsorted(cdf, value.numpy() if isinstance(value, ops.Tensor) else value, side=side) + expected_value = value + if isinstance(value, ops.Tensor): + expected_value = self.evaluate(value) + result = np.searchsorted(cdf, expected_value, side=side) tf_result = self.evaluate( array_ops.searchsorted(cdf, value, side=side)) self.assertAllEqual(result, tf_result) From b71f27a9f91e4f1d9fc01b17fdffa67642964f9e Mon Sep 17 00:00:00 2001 From: Khush Date: Tue, 22 Sep 2026 23:54:45 +0530 Subject: [PATCH 16/35] test: fix scalar searchsorted test type check --- tensorflow/python/kernel_tests/array_ops/array_ops_test.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 4ff7180fc7c103..6f722be1123b22 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -2111,10 +2111,7 @@ def testScalarValue(self): for value in values: for side in ("left", "right"): with self.subTest(value=value, side=side): - expected_value = value - if isinstance(value, ops.Tensor): - expected_value = self.evaluate(value) - result = np.searchsorted(cdf, expected_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) From d95c94d26c4e0f05aa0e58518d5f3605f48c0d32 Mon Sep 17 00:00:00 2001 From: Akhil Goel Date: Wed, 23 Sep 2026 00:31:25 -0700 Subject: [PATCH 17/35] PR #49301: [XLA:GPU][oneAPI] Allow null pointer kernel arguments for sycl kernels Imported from GitHub PR https://github.com/openxla/xla/pull/49301 LLVM might retain unused 0-sized kernel arguments and mark them as `readnone` instead of removing them. The runtime will not allocate any memory and return nullptr. Hence, a valid kernel can carry a null pointer argument. This PR relaxes the sycl kernel launch null argument check so that valid kernels with nullptr arguments are not rejected. This helps fix `xla/backends/gpu/codegen/emitters/tests/scatter/sorted_empty_indices.hlo` test with the oneAPI backend. This issue came to light after https://github.com/openxla/xla/commit/f12e3c4264471c510649e30533185878688861fa added a test with empty indices for a scatter operation. Copybara import of the project: -- db19166bd5dbbe92137acc7ecd33013fec588af0 by Akhil Goel : Allow nullptr kernel args Merging this change closes #49301 PiperOrigin-RevId: 986541498 --- third_party/xla/xla/stream_executor/sycl/sycl_stream.cc | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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]); From 9a4cf0f33848ceadd75edb9a2d55b1be10edb307 Mon Sep 17 00:00:00 2001 From: Om Thakkar Date: Wed, 23 Sep 2026 00:49:08 -0700 Subject: [PATCH 18/35] PR #48307: [XLA:CPU] Explicitly disable oneDNN for YNN-specific tests Imported from GitHub PR https://github.com/openxla/xla/pull/48307 This PR explicitly disables oneDNN custom calls for YNN-specific tests to avoid unit test failures when running with oneDNN (as `--xla_cpu_experimental_onednn_custom_call` is `false` by default). Copybara import of the project: -- 9df06a9bfb4082b5adb3e3623638ed9d5909fb92 by Om Thakkar : explicitly disable oneDNN for YNN-specific tests Merging this change closes #48307 PiperOrigin-RevId: 986549795 --- .../xla/xla/backends/cpu/transforms/ynn_matcher_test.cc | 4 ++++ third_party/xla/xla/service/cpu/cpu_compiler_test.cc | 1 + 2 files changed, 5 insertions(+) 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/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); From e74627317123379b094622f18644474b90f632f1 Mon Sep 17 00:00:00 2001 From: gaurides Date: Wed, 23 Sep 2026 01:04:45 -0700 Subject: [PATCH 19/35] PR #43318: [oneDNN][CPU] Add initial FP8 support for oneDNN dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/43318 ๐Ÿ“ Summary of Changes Introduce AMX-FP8 CPU feature and use it to gate FP8 dispatch through oneDNN, with an AMX-FP16 emulation fallback for hosts without native AMX-FP8. - Add CPUFeature::AMX_FP8, map it to `diamondrapids`, and register it in CpuFeatureFromString. - Detect `+amx-fp8` / `+amx-fp16` from the LLVM feature string and expose has_amx_fp8()/has_amx_fp16() on TargetMachineFeatures. - When max_cpu_feature is capped at AMX_FP16, suppress DMR-and-newer additions (amx-fp8/-tf32/-avx512/-movrs, avx10.2, APX) so a GNR-tuned binary cannot emit post-GNR instructions. - Extend IsOneDnnSupportedDType to accept F8E5M2/F8E4M3FN/F8E4M3: dispatch natively on AMX-FP8, else fall back to AMX-FP16 emulation and LOG_FIRST_N the fallback. - Add a parameterized TargetMachineFeatures test covering SPR, GNR (emulation), and DMR (native) feature strings. ๐ŸŽฏ Justification We need to detect if x86 CPU supports FP8 because calling onednn for fp8 functionality ๐Ÿš€ Kind of Contribution : โœจ New Feature, ๐Ÿ“Š Benchmark (for Performance Improvements) This doesn't affect any performance yet. ๐Ÿงช Unit Tests: xla/service/cpu/tests/cpu_vectorization_test.cc ๐Ÿงช Execution Tests: Doesn't affect existing execution tests Copybara import of the project: -- 7d975a8af8825ac9afcde6e6b9b9a0788bb112ee by Gauri Deshpande : [oneDNN][CPU] Add initial FP8 support for oneDNN dispatch -- 9445a7dad3b751228a6a7f4f23383fee05eadec5 by Gauri Deshpande : Fix clang-tidy braces and missing absl/log include in FP8 oneDNN support -- b33af61a351de16f8706e6b0ed26a2bcc0a95f16 by gaurides : Address review comment - use AVX10.2 Co-authored-by: Penporn Koanantakool <38085909+penpornk@users.noreply.github.com> -- 60727426128036d9926dbca6a7efd29087ec5017 by Gauri Deshpande : fix rebase duplicate lines of codes Merging this change closes #43318 PiperOrigin-RevId: 986557057 --- .../xla/backends/cpu/codegen/cpu_features.cc | 16 +++++++ .../cpu/codegen/target_machine_features.cc | 10 ++-- .../cpu/codegen/target_machine_features.h | 4 ++ .../codegen/target_machine_features_test.cc | 46 +++++++++++++++++++ .../xla/xla/backends/cpu/onednn_support.cc | 26 +++++++++-- third_party/xla/xla/debug_options_flags.cc | 4 +- .../xla/xla/debug_options_parsers_test.cc | 1 + third_party/xla/xla/service/cpu/onednn_util.h | 14 ++++++ .../cpu/tests/cpu_vectorization_test.cc | 20 ++++++++ 9 files changed, 132 insertions(+), 9 deletions(-) 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/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index d2db95a8c0da98..cd98e0f59b26a6 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -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/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}, }); } From 454d52bd10fda927420dd27b01ea1883f9f8bdd9 Mon Sep 17 00:00:00 2001 From: dlcompilers-infra-bot Date: Wed, 23 Sep 2026 01:04:46 -0700 Subject: [PATCH 20/35] PR #48288: add parallel threads to explicitly annotated instructions so LHS can treat it as async ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48288 ๐Ÿ“ Summary of Changes Add the "parallel" to the execution threads for any explicitly stream-annotated instructions. LHS checks for the execution threads to determine if a certain computation contains async ops for scheduling, without this, if a computation only contains stream-annotated instructions, LHS will skip scheduling it. ๐ŸŽฏ Justification functional correctness for stream-annotated instructions ๐Ÿš€ Kind of Contribution Please remove what does not apply: ๐Ÿ› Bug Fix ๐Ÿงช Unit Tests: added unit test Copybara import of the project: -- eb5031dff696996235c405ec1ca7bc564d6198d0 by TJ Xu : add parallel threads to explicitly annotated instructions so LHS can treat it as async ops -- 60281ae08a36d6035b6b9d8758e9a1162926cb41 by TJ Xu : address pr comments -- 0a65afaf3b674990b29553a94fafb547b7d6f7fe by TJ Xu : fix test failure -- 0e182bdd1fdf9f7fcc5eeeef038c64dd7bfd9aa3 by TJ Xu : fix decomposer when called by autotuner Merging this change closes #48288 PiperOrigin-RevId: 986557067 --- ...xplicit_stream_annotation_async_wrapper.cc | 9 ++- ...explicit_stream_annotation_async_wrapper.h | 2 - ...it_stream_annotation_async_wrapper_test.cc | 56 +++++++++++++-- .../stream_attribute_async_wrapper.cc | 9 ++- .../stream_attribute_async_wrapper.h | 2 - .../codegen/emitters/fusion_wrapper_base.cc | 2 + third_party/xla/xla/hlo/ir/hlo_instruction.h | 1 + third_party/xla/xla/service/gpu/BUILD | 1 - .../xla/xla/service/gpu/gpu_hlo_schedule.cc | 3 +- third_party/xla/xla/tools/hlo_decomposer.cc | 20 +++++- .../xla/xla/tools/hlo_decomposer_test.cc | 69 +++++++++++++++++++ 11 files changed, 150 insertions(+), 24 deletions(-) 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/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/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/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/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/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"( From 708998581eb82269da3c0284661e5d828f8c6e93 Mon Sep 17 00:00:00 2001 From: Aliia Khasanova Date: Wed, 23 Sep 2026 02:05:47 -0700 Subject: [PATCH 21/35] Add kUpdateGolden mode for AOT compatibility goldens Goldens could only be written by hand. kUpdateGolden dumps them from a test run; the script that drives it follows in the next change. Two behaviour changes: * A missing target directory is now NotFound rather than an empty test matrix. A mistyped target name used to yield a suite that passed while testing nothing. * AOTTestPlatform is required on the test_lib entry points. The GPU default let a CPU caller resolve GPU paths. Dumping and verification must compile under identical debug options or the golden is not comparable to what verification recompiles, so `ApplyAotDeterminismDebugOptions` is the single definition and every CompileOptions site routes through it. PiperOrigin-RevId: 986588877 --- third_party/xla/xla/tests/BUILD | 2 + .../aot_compatibility_experimental/BUILD | 6 + .../aot_compatibility_experimental/gpu/BUILD | 1 + .../gpu/collective_ops_aot_test.cc | 7 +- .../executables/test_dummy_test/v1/exec.pbtxt | 7 + .../test_lib.cc | 142 +++++--- .../aot_compatibility_experimental/test_lib.h | 51 +-- .../test_lib_test.cc | 272 ++++++++++++++- .../xla/tests/aot_interception_pjrt_client.cc | 322 ++++++++++-------- .../xla/tests/aot_interception_pjrt_client.h | 14 +- 10 files changed, 601 insertions(+), 223 deletions(-) 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_; From 22a4e6f331b14c7780f3944b344d2ad899d7f0d6 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Wed, 23 Sep 2026 02:23:22 -0700 Subject: [PATCH 22/35] In OpenXLA builds, `googletest.patch` adds `hdrs` to `cc_library(name = "gtest_main")` for Bazel `layering_check` compatibility, which disables BANT DWYU's empty-`hdrs` keep heuristic. Without `includes` or `keep_dep` on `gtest_main`, BANT DWYU fails to index `gtest_main`'s headers under `gtest/*.h` and `gmock/*.h` and suggests removing `@com_google_googletest//:gtest_main` from test targets. Add `includes` and `tags = ["avoid_dep", "keep_dep"]` to `gtest_main` in `googletest.patch` so `gtest_main` satisfies `gtest`/`gmock` header includes and is never removed by DWYU, while `avoid_dep` ensures `:gtest` is still preferred when suggesting missing dependencies (and omitting `alwayslink = True` avoids duplicate `main()` ODR violations in targets that also transitively depend on `//tensorflow/core:test_main`). PiperOrigin-RevId: 986597415 --- .../xla/build_tools/lint/run_dwyu_test.py | 108 ++++++++++++++++++ .../third_party/googletest/googletest.patch | 16 ++- 2 files changed, 123 insertions(+), 1 deletion(-) 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": [], From c6e859b56bfc441cc1adcd34068545adfb72cc4a Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 02:37:42 -0700 Subject: [PATCH 23/35] [XLA:CPU] Benchmarks: Disable HLO modules upload. PiperOrigin-RevId: 986605327 --- .../xla/xla/backends/cpu/benchmarks/hlo_benchmark_runner.cc | 3 +++ 1 file changed, 3 insertions(+) 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) { From fd275d220665bb7236014d300cf39f1ac556c3f0 Mon Sep 17 00:00:00 2001 From: Mohammed Anany Date: Wed, 23 Sep 2026 03:14:16 -0700 Subject: [PATCH 24/35] [XLA:GPU][NFC] Fix pre-existing clang-tidy findings in the GPU autotuner and cost model PiperOrigin-RevId: 986624882 --- .../xla/xla/backends/gpu/autotuner/BUILD | 12 +++--- .../backends/gpu/autotuner/autotuner_main.cc | 1 - .../gpu/autotuner/block_level_emitter.cc | 3 -- .../gpu/autotuner/block_level_emitter_test.cc | 20 ++++++---- .../xla/xla/backends/gpu/autotuner/factory.h | 1 + .../backends/gpu/autotuner/factory_test.cc | 39 ++++++++++++++++--- .../xla/xla/backends/gpu/transforms/BUILD | 1 - .../softmax_rewriter_triton_test.cc | 1 - third_party/xla/xla/service/gpu/model/BUILD | 2 +- .../gpu_indexing_performance_model_test.cc | 2 +- 10 files changed, 54 insertions(+), 28 deletions(-) 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/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/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/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" From 520228a81f19bf1b95f3958e5278ac53e17b47dd Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 03:29:40 -0700 Subject: [PATCH 25/35] Automated Code Change PiperOrigin-RevId: 986632786 --- tensorflow/lite/delegates/flex/kernel.cc | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) 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 { From ebbf121b13a5a0f26e679497d9899ec318481fe2 Mon Sep 17 00:00:00 2001 From: Mikhail Goncharov Date: Wed, 23 Sep 2026 03:59:04 -0700 Subject: [PATCH 26/35] [XLA:GPU] switch tiling block-level emitter on by default This changes tiling analysis and block level emitter. While we have tried to keep the implementation close to current triton block-level emitter there are some differences that migth result in: - different numeric results - different fusion decisions, due to cost model estimations - new bugs - increased compilation time - mostly due to priority fusion. We are working on optimizing it. PiperOrigin-RevId: 986646237 --- third_party/xla/xla/debug_options_flags.cc | 2 +- third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index cd98e0f59b26a6..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 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, From 6a04fe8663b5bcc442c99f1cb069c969b9d426ae Mon Sep 17 00:00:00 2001 From: Alexandros Theodoridis Date: Wed, 23 Sep 2026 04:28:40 -0700 Subject: [PATCH 27/35] PR #49152: [ROCm] Switch to newest rules_ml_toolchain, override rocm_hermetic_download MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49152 ๐Ÿ“ Summary of Changes Switch to newest rules_ml_toolchain, override hermetic rocm download in bzlmod ๐ŸŽฏ Justification In a past implementation we used hermetic_rocm_download from inside rules_ml_toolchain that leads to a usage of the older rocm distribution for the rocm_toolchain (hipcc). In this PR we override the extension hence using whatever is provided by the rocm_configure. ๐Ÿš€ Kind of Contribution Please remove what does not apply: ๐Ÿ› Bug Fix, ๐Ÿ“Š Benchmark (for Performance Improvements) Not relevant ๐Ÿงช Unit Tests: CI ๐Ÿงช Execution Tests: CI Copybara import of the project: -- 2e7c9dffb75b39a2b737a1d6707e3b30494db9bb by Alexandros Theodoridis : Switch to newest rules_ml_toolchain, override rocm_hermetic_download -- 9a2fb9b8f25f523ac3c2884e625dc15e85fe4767 by Alexandros Theodoridis : Override rocm_hermetic_download_ext -- e83ec060a34220c2424f090a8d74e1fee285fc75 by Alexandros Theodoridis : Fix invalid hipcc usage Merging this change closes #49152 PiperOrigin-RevId: 986660810 --- third_party/xla/MODULE.bazel | 13 ++++++++++--- third_party/xla/workspace3.bzl | 6 +++--- 2 files changed, 13 insertions(+), 6 deletions(-) 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/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", ), ) From f8316770739fe10b9fa49fb3f90223ef1c389263 Mon Sep 17 00:00:00 2001 From: Misha Gutman Date: Wed, 23 Sep 2026 04:51:22 -0700 Subject: [PATCH 28/35] Add quantized (INT8/INT4/INT2 weights) and FP16/BF16 execution to the YNNPACK delegate's MoE op. PiperOrigin-RevId: 986670634 --- tensorflow/lite/delegates/ynnpack/BUILD | 3 + tensorflow/lite/delegates/ynnpack/moe.cc | 445 ++++++++--- .../lite/delegates/ynnpack/moe_bench.cc | 739 ++++++++++++++++-- tensorflow/lite/delegates/ynnpack/moe_test.cc | 688 +++++++++++++++- 4 files changed, 1724 insertions(+), 151 deletions(-) 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 From d1456a9d9b25ae335f1d98acfb3627ad08ce5bd5 Mon Sep 17 00:00:00 2001 From: Enver Kayaaslan Date: Wed, 23 Sep 2026 04:52:55 -0700 Subject: [PATCH 29/35] Pass specialisation cache directly on ctor and drop is_embedding bool on hlo evaluators. Also drop clearing the cache for computations and instructions on child evaluators. Using caching does not support for when the computations may be modified. PiperOrigin-RevId: 986671239 --- .../xla/xla/hlo/evaluator/hlo_evaluator.cc | 26 ++---- .../xla/xla/hlo/evaluator/hlo_evaluator.h | 46 +++++----- .../xla/hlo/evaluator/hlo_evaluator_test.cc | 87 ++++++++++++------- 3 files changed, 89 insertions(+), 70 deletions(-) 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) { From c5ac86dcf0bd1393057de8b7d1f9ef76579df7cc Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 04:54:28 -0700 Subject: [PATCH 30/35] Automated Code Change PiperOrigin-RevId: 986671871 --- third_party/xla/xla/codegen/xtile/codegen/BUILD | 1 + .../xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc | 1 + 2 files changed, 2 insertions(+) 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 { From d36a847c5c8d3c816abff309c90a4ff7f6a17f38 Mon Sep 17 00:00:00 2001 From: Sohaib Iftikhar Date: Wed, 23 Sep 2026 05:04:44 -0700 Subject: [PATCH 31/35] [XLA:GPU]: Allow CUFunction in Record API Record runs after Prepare -> Initialize. The following events can happen given 2 GPUs within the same process. TS1. Thread0 (Host) queues a collective kernel on Gpu0. Asynchronously. Gpu0 kernel (device) blocks on GPU1 to also execute the corresponding collective on Gpu1. TS2. Thread0 Records a kernal (with FFI). Calls cuModuleLoad and acquires driver mutex. However loading cannot finish until Gpu0 (device) is free. So it waits. TS3. Thread1 which lags a bit behind Thread0 calls cuLaunchKernel but blocks because the driver lock was taken during TS2. Collective kernel is never launched and GPU0 (TS1) keeps waiting. To prevent this deadlock we must allow FFI users to load kernels during preparation. Since not all FFI users can use se::Kernel API for loading we can allow using a function pointer that was loaded and cached by the FFI user so cuModuleLoad is not called anymore during the record phase. Note that this problem only exists for users running single process across multiple GPUs. PiperOrigin-RevId: 986676717 --- third_party/xla/xla/ffi/api/record_api.h | 16 +- third_party/xla/xla/ffi/api/record_c_api.h | 11 +- .../xla/xla/ffi/api/record_ffi_test.cc | 122 +++++++++++++++- third_party/xla/xla/ffi/record_ffi_test.cc | 138 ++++++++++++++++-- 4 files changed, 255 insertions(+), 32 deletions(-) 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 From 563b176b6af7dac4817a1ae3576aa529e778862b Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 23 Sep 2026 05:29:30 -0700 Subject: [PATCH 32/35] [XLA:GPU] Remove side-effecting heroes in DynamicSliceFusionRewriterV2 When `RewriteHero` replaces downstream `dynamic-update-slice` or `get-tuple-element` instructions rather than the `hero` instruction directly, `HloComputation::ReplaceInstruction` invokes `RemoveInstructionAndUnusedOperands` on the downstream instruction. Because `RemoveInstructionAndUnusedOperands` skips upstream operands that have side effects (`item->HasSideEffect() && item != instruction`), a side-effecting hero (such as an FFI custom-call with `custom_call_has_side_effect=true`) and its sliced operand chain remain alive in the parent computation alongside the newly created `dynamic_slice_fusion`, causing the side-effecting custom-call to execute twice. Transfer control dependencies from `hero` to `fusion` and explicitly invoke `RemoveInstructionAndUnusedOperands(hero)` when `hero` has side effects and is not directly replaced. Also check user counts in `ResolveLeafDus` and `ResolveNonTupleSlicedResult` (and ensure tuple heroes only have unique GTE users in `HasSupportedShapes`) so DUS chains are only absorbed when the hero/GTE has a single user, preventing duplicate heroes when external users exist. PiperOrigin-RevId: 986687589 --- .../dynamic_slice_fusion_rewriter_v2.cc | 38 +++ .../dynamic_slice_fusion_rewriter_v2_test.cc | 261 ++++++++++++++++++ 2 files changed, 299 insertions(+) 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 From 45909285862ead9e72205148a85ea27fd8aa7972 Mon Sep 17 00:00:00 2001 From: Amit Sabne Date: Wed, 23 Sep 2026 05:54:02 -0700 Subject: [PATCH 33/35] Remove unused memory_space_assignment implementation from open-source XLA. The `xla::memory_space_assignment` pass implementation, options, simulator, cost analysis, and repacking algorithms are only used by internal backend compilation pipelines, whereas open-source `buffer_assignment.h` only references `xla::memory_space_assignment::PresetAssignments`. Remove the unused implementation files and unit tests from `xla/service/memory_space_assignment/`, retaining only `PresetAssignments` in `xla/service/memory_space_assignment/memory_space_assignment.h`. PiperOrigin-RevId: 986698639 --- tensorflow/core/BUILD | 1 - .../xla/service/memory_space_assignment/BUILD | 731 +- .../memory_space_assignment/algorithm.cc | 11102 --------- .../memory_space_assignment/algorithm.h | 1956 -- .../memory_space_assignment/allocation.cc | 1402 -- .../memory_space_assignment/allocation.h | 736 - .../allocation_test.cc | 300 - .../allocation_value.cc | 112 - .../allocation_value.h | 342 - .../best_fit_repacker.cc | 650 - .../best_fit_repacker.h | 77 - .../best_fit_repacker_test.cc | 510 - .../buffer_interval_comparator.cc | 172 - .../buffer_interval_comparator.h | 150 - .../memory_space_assignment/cost_analysis.cc | 502 - .../memory_space_assignment/cost_analysis.h | 268 - .../cost_analysis_test.cc | 297 - .../live_range_util.cc | 511 - .../memory_space_assignment/live_range_util.h | 241 - .../live_range_util_test.cc | 795 - .../memory_bound_loop_optimizer.cc | 1630 -- .../memory_bound_loop_optimizer.h | 484 - .../memory_bound_loop_optimizer_test.cc | 1746 -- .../memory_space_assignment.cc | 1713 -- .../memory_space_assignment.h | 171 - .../memory_space_assignment.proto | 261 - .../memory_space_assignment_test.cc | 20049 ---------------- .../memory_space_assignment_test_base.h | 634 - .../memory_space_assignment/options.cc | 139 - .../service/memory_space_assignment/options.h | 504 - .../prefetch_interval_picker.cc | 587 - .../prefetch_interval_picker.h | 297 - .../prefetch_interval_picker_test.cc | 453 - .../memory_space_assignment/repacking.h | 49 - .../memory_space_assignment/simulator.cc | 450 - .../memory_space_assignment/simulator.h | 184 - .../memory_space_assignment/simulator_test.cc | 636 - .../service/memory_space_assignment/slice.cc | 160 - .../service/memory_space_assignment/slice.h | 150 - .../memory_space_assignment/slice_test.cc | 224 - .../memory_space_assignment/testing_utils.h | 139 - .../memory_space_assignment/tuning_utils.cc | 64 - .../memory_space_assignment/tuning_utils.h | 43 - .../service/memory_space_assignment/utils.cc | 607 - .../service/memory_space_assignment/utils.h | 165 - third_party/xla/xla/xla.default.bzl | 1 - 46 files changed, 3 insertions(+), 52392 deletions(-) delete mode 100644 third_party/xla/xla/service/memory_space_assignment/algorithm.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/algorithm.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation_value.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/allocation_value.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/best_fit_repacker.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/best_fit_repacker_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/buffer_interval_comparator.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/cost_analysis.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/cost_analysis.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/cost_analysis_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/live_range_util.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/live_range_util.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/live_range_util_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_bound_loop_optimizer_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/options.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/options.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/repacking.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/simulator.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/simulator.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/simulator_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/slice.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/slice.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/slice_test.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/testing_utils.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/tuning_utils.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/tuning_utils.h delete mode 100644 third_party/xla/xla/service/memory_space_assignment/utils.cc delete mode 100644 third_party/xla/xla/service/memory_space_assignment/utils.h 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/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/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"), From e89f730e3b14bd69dc8ed08598af1c91ec46a4a5 Mon Sep 17 00:00:00 2001 From: Viktor Moros Date: Wed, 23 Sep 2026 06:03:21 -0700 Subject: [PATCH 34/35] [XLA Highway] Add block interleaving & lane permutation helpers PiperOrigin-RevId: 986702595 --- third_party/xla/xla/pjrt/transpose_kernels.h | 54 ++++++++++++++++++++ 1 file changed, 54 insertions(+) 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(); From bf6acbb271b7f41de787e7bf485fcd1397400d1a Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 23 Sep 2026 06:04:14 -0700 Subject: [PATCH 35/35] Automated Code Change PiperOrigin-RevId: 986703064 --- .../xla/backends/gpu/codegen/cudnn_test.cc | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) 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);