From cfa54990ca3c27bcb04c2992668a063b7d0f5948 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Sun, 26 Apr 2026 23:23:28 +0530 Subject: [PATCH 01/25] Fix tf.math.sign returning 0 for small nonzero complex numbers Eigen's scalar_sign_op computes |z|^2 = re^2 + im^2 as a float32 intermediate. When |z| < sqrt(float32_tiny) ~= 1.08e-19, this underflows to 0, causing sign(z) to incorrectly return 0 instead of z/|z|. Add specialized sign functors for complex64 and complex128 that check (re != 0 || im != 0) directly and use std::abs() (which uses hypot internally) to compute the magnitude without intermediate underflow/overflow. Fixes #116945 --- tensorflow/core/kernels/cwise_ops.h | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h index 06d753722e7c9a..33d641c5301c24 100644 --- a/tensorflow/core/kernels/cwise_ops.h +++ b/tensorflow/core/kernels/cwise_ops.h @@ -865,6 +865,32 @@ struct log1p : base> {}; template struct sign : base> {}; +// Specialization of sign for complex types to avoid underflow in |z|^2. +// Eigen's scalar_sign_op computes |z|^2 = re^2 + im^2 as a float32 +// intermediate. When |z| < sqrt(float32_tiny) ~= 1.08e-19, this underflows +// to 0, causing sign(z) to incorrectly return 0. We fix this by checking +// whether the input is nonzero directly, and using std::abs() (which uses +// hypot internally) to avoid intermediate underflow/overflow. +template +struct safe_complex_sign_op { + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex operator()( + const std::complex& z) const { + if (z.real() == RealType(0) && z.imag() == RealType(0)) { + return std::complex(0, 0); + } + return z / std::abs(z); + } +}; + +template <> +struct sign + : base> {}; + +template <> +struct sign + : base> {}; + + template struct sinh : base> {}; From 34b646bf922e30204de736b8e04d62198ca5c413 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Mon, 27 Apr 2026 20:43:22 +0530 Subject: [PATCH 02/25] Add regression test for complex sign with small magnitudes Add testComplexSignSmallMagnitude to cwise_ops_unary_test.py that verifies tf.math.sign returns correct unit-magnitude results for complex64 and complex128 inputs with magnitudes below sqrt(float32_tiny), where the original |z|^2 computation would underflow. Requested by reviewer in PR #116986. --- .../math_ops/cwise_ops_unary_test.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py index 24c06cedce2443..f4318c302ec407 100644 --- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py +++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py @@ -651,6 +651,37 @@ def g(x): *gradient_checker_v2.compute_gradient(g, [ops.convert_to_tensor(2.0)])) self.assertLess(err, 1e-3) + @test_util.run_in_graph_and_eager_modes + def testComplexSignSmallMagnitude(self): + """Regression test for #116945: sign(z) must not return 0 for small |z|. + + Eigen's scalar_sign_op computes |z|^2 = re^2 + im^2 as a float32 + intermediate, which underflows to 0 when |z| < sqrt(float32_tiny) ~= + 1.08e-19, causing sign(z) to incorrectly return 0. + """ + # Test values with magnitudes below the underflow threshold. + for dtype in [np.complex64, np.complex128]: + small_values = np.array([ + 1e-20 + 0j, + 1e-30 + 0j, + 0 + 1e-20j, + -1e-20 + 0j, + 1e-20 + 1e-20j, + ], dtype=dtype) + + result = self.evaluate(math_ops.sign(constant_op.constant(small_values))) + + # sign(z) for nonzero z must be z / |z| (unit magnitude), never 0. + for i, z in enumerate(small_values): + expected = z / np.abs(z) + self.assertNotEqual( + result[i], 0, + msg=f"sign({z}) returned 0 for dtype={dtype.__name__}") + self.assertAllClose( + result[i], expected, rtol=1e-5, + msg=f"sign({z}) incorrect for dtype={dtype.__name__}") + + if __name__ == "__main__": test.main() From c074672e99e7f98910ac181b5352a0820d9bc115 Mon Sep 17 00:00:00 2001 From: Sachin Mahajan Date: Wed, 1 Jul 2026 01:57:36 +0530 Subject: [PATCH 03/25] Use std::hypot instead of std::abs to fix CPU underflow in complex sign std::abs(complex) may compute sqrt(re*re + im*im) which underflows to 0 when FTZ is enabled and |re|^2 is subnormal. Using std::hypot directly avoids intermediate underflow/overflow per C++ standard guarantees. Also avoids complex division in favor of scalar re/mag, im/mag for numerical robustness. --- tensorflow/core/kernels/cwise_ops.h | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h index 33d641c5301c24..8c30994a25d727 100644 --- a/tensorflow/core/kernels/cwise_ops.h +++ b/tensorflow/core/kernels/cwise_ops.h @@ -875,10 +875,16 @@ template struct safe_complex_sign_op { EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex operator()( const std::complex& z) const { - if (z.real() == RealType(0) && z.imag() == RealType(0)) { + const RealType re = z.real(); + const RealType im = z.imag(); + if (re == RealType(0) && im == RealType(0)) { return std::complex(0, 0); } - return z / std::abs(z); + // Use hypot directly to avoid intermediate underflow/overflow. + // std::abs(complex) may compute sqrt(re*re + im*im) which underflows + // when FTZ is enabled and |re|^2 or |im|^2 is subnormal. + const RealType mag = std::hypot(re, im); + return std::complex(re / mag, im / mag); } }; From b4c65edd29d94fa73219a1ef6a8e7d5444816fb2 Mon Sep 17 00:00:00 2001 From: adi-IL Date: Sun, 23 Aug 2026 03:58:15 +0530 Subject: [PATCH 04/25] fix: validate filter and out_backprop dimensions in Conv2DBackpropInput to prevent crash on invalid rank --- tensorflow/core/kernels/conv_grad_input_ops.h | 12 ++++++++-- .../core/kernels/conv_grad_shape_utils.cc | 11 ++++++++++ .../nn_ops/conv2d_transpose_test.py | 22 +++++++++++++++++++ tensorflow/python/ops/nn_ops.py | 5 +++++ 4 files changed, 48 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/kernels/conv_grad_input_ops.h b/tensorflow/core/kernels/conv_grad_input_ops.h index d28077272041b0..9140a987b23bf6 100644 --- a/tensorflow/core/kernels/conv_grad_input_ops.h +++ b/tensorflow/core/kernels/conv_grad_input_ops.h @@ -356,7 +356,11 @@ class Conv2DBackpropInputOp : public OpKernel { OP_REQUIRES( context, out_backprop.dims() == 4, absl::InvalidArgumentError(absl::StrCat( - "input_sizes must be 4-dimensional, got: ", out_backprop.dims()))); + "out_backprop must be 4-dimensional, got: ", out_backprop.dims()))); + OP_REQUIRES( + context, filter.dims() == 4, + absl::InvalidArgumentError(absl::StrCat( + "filter must be 4-dimensional, got: ", filter.dims()))); TensorShape input_shape; OP_REQUIRES_OK(context, @@ -468,7 +472,11 @@ class Conv2DCustomBackpropInputOp : public OpKernel { OP_REQUIRES( context, out_backprop.dims() == 4, absl::InvalidArgumentError(absl::StrCat( - "input_sizes must be 4-dimensional, got: ", out_backprop.dims()))); + "out_backprop must be 4-dimensional, got: ", out_backprop.dims()))); + OP_REQUIRES( + context, filter.dims() == 4, + absl::InvalidArgumentError(absl::StrCat( + "filter must be 4-dimensional, got: ", filter.dims()))); TensorShape input_shape; OP_REQUIRES_OK(context, diff --git a/tensorflow/core/kernels/conv_grad_shape_utils.cc b/tensorflow/core/kernels/conv_grad_shape_utils.cc index 3db06801a9a6fa..fbace3241e0d60 100644 --- a/tensorflow/core/kernels/conv_grad_shape_utils.cc +++ b/tensorflow/core/kernels/conv_grad_shape_utils.cc @@ -174,6 +174,17 @@ absl::Status Conv2DBackpropComputeInputShape( const Tensor& input_sizes, const TensorShape& filter_shape, const TensorShape& out_backprop_shape, const TensorFormat& data_format, TensorShape* input_shape) { + if (filter_shape.dims() != 4) { + return absl::InvalidArgumentError(absl::StrCat( + "Conv2DBackpropInput: filter must be 4-dimensional, not ", + filter_shape.dims())); + } + if (out_backprop_shape.dims() != 4) { + return absl::InvalidArgumentError(absl::StrCat( + "Conv2DBackpropInput: out_backprop must be 4-dimensional, not ", + out_backprop_shape.dims())); + } + if (!TensorShapeUtils::IsVector(input_sizes.shape())) { return absl::InvalidArgumentError(absl::StrCat( "Conv2DBackpropInput: input_sizes input must be 1-dim, not ", diff --git a/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py b/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py index 98087520bd492b..8fb56ba136c875 100644 --- a/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/conv2d_transpose_test.py @@ -332,5 +332,27 @@ def testConv2DTransposeLargeOutputShape(self): filters=[[[[1]]]]) self.evaluate(op) + def testConv2DTransposeInvalidFilterRank(self): + with self.session(): + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + op = nn_ops.conv2d_transpose( + input=constant_op.constant(1.0, shape=[1, 3, 3, 8]), + filters=8, + output_shape=(2, 2), + strides=(2, 2), + padding="SAME") + self.evaluate(op) + + def testConv2DTransposeInvalidOutputShapeRank(self): + with self.session(): + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + op = nn_ops.conv2d_transpose( + input=constant_op.constant(1.0, shape=[1, 4, 4, 3]), + filters=constant_op.constant(1.0, shape=[3, 3, 3, 3]), + output_shape=(2, 4, 2), + strides=[1, 2, 2, 1], + padding="SAME") + self.evaluate(op) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py index 2948e23a910dfb..a041f8e5026f10 100644 --- a/tensorflow/python/ops/nn_ops.py +++ b/tensorflow/python/ops/nn_ops.py @@ -2754,6 +2754,11 @@ def conv2d_transpose_v2( data_format = "NHWC" channel_index = 1 if data_format.startswith("NC") else 3 + filters = ops.convert_to_tensor(filters, name="filters") + if filters.shape.rank is not None and filters.shape.rank != 4: + raise ValueError(f"`filters` must be 4-dimensional. " + f"Received: filters with shape {filters.shape}") + strides = _get_sequence(strides, 2, channel_index, "strides") dilations = _get_sequence(dilations, 2, channel_index, "dilations") padding, explicit_paddings = convert_padding(padding) From beafd143cefd1c3e1a08ed369bb1470132050b5b Mon Sep 17 00:00:00 2001 From: kokol16 Date: Fri, 4 Sep 2026 16:54:01 -0400 Subject: [PATCH 05/25] Fix GPU tiled-transpose index overflow The tiled-transpose kernel used 32-bit indexes, which overflowed for tensors with more than INT32_MAX elements and caused out-of-bounds memory access. Use 64-bit indexes in that kernel and add a test for an index above INT32_MAX. --- tensorflow/core/kernels/BUILD | 12 ++++++ tensorflow/core/kernels/conv_2d_gpu.h | 26 +++++++------ .../core/kernels/conv_2d_gpu_test.cu.cc | 37 +++++++++++++++++++ 3 files changed, 64 insertions(+), 11 deletions(-) create mode 100644 tensorflow/core/kernels/conv_2d_gpu_test.cu.cc diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index 39eb110d9c6d91..c33068afa5a53f 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -335,6 +335,18 @@ tf_kernel_library( alwayslink = 1, ) +tf_cuda_only_cc_test( + name = "conv_2d_gpu_test", + size = "small", + srcs = ["conv_2d_gpu_test.cu.cc"], + features = if_cuda(["-layering_check"]), + deps = [ + ":conv_2d", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + cc_library( name = "conv_2d_hdrs", hdrs = ["conv_2d.h"], diff --git a/tensorflow/core/kernels/conv_2d_gpu.h b/tensorflow/core/kernels/conv_2d_gpu.h index 3c9ccb5e6330f8..89e53268c814a1 100644 --- a/tensorflow/core/kernels/conv_2d_gpu.h +++ b/tensorflow/core/kernels/conv_2d_gpu.h @@ -150,10 +150,10 @@ struct Index : Array { }; // A helper function that converts a tensor index into a flat array index. -template -EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int TensorIndexToFlat( +template +EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE FlatIndex TensorIndexToFlat( const Index& index, const Dimension& dims) { - int flat_index = index[0]; + FlatIndex flat_index = index[0]; for (int i = 1; i < IndexCount; i++) { flat_index = flat_index * dims[i] + index[i]; } @@ -320,8 +320,8 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( input_tile_index[2] * TileSizeJ, }; - int input_origin_flat_index = - TensorIndexToFlat(input_tile_origin, input_dims); + int64_t input_origin_flat_index = + TensorIndexToFlat(input_tile_origin, input_dims); bool full_tile = true; int tile_width = TileSizeJ; @@ -350,8 +350,10 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( // dimension of the input array. int ti = x / TileSizeJ; int tj = x % TileSizeJ; - int input_index = input_origin_flat_index + ti * input_dims[2] + tj; - int input_increment = ReadRowPerPass * input_dims[2]; + int64_t input_index = + input_origin_flat_index + static_cast(ti) * input_dims[2] + tj; + int64_t input_increment = + static_cast(ReadRowPerPass) * input_dims[2]; if (full_tile) { #pragma unroll @@ -385,8 +387,8 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( output_tile_index[2] * TileSizeI, }; - int output_origin_flat_index = - TensorIndexToFlat(output_tile_origin, output_dims); + int64_t output_origin_flat_index = + TensorIndexToFlat(output_tile_origin, output_dims); constexpr int out_effective_thread_num = NumThreads / TileSizeI * TileSizeI; @@ -396,8 +398,10 @@ __global__ void SwapDimension1And2InTensor3UsingTiles( // dimension of the output array. int ti = x / TileSizeI; int tj = x % TileSizeI; - int output_index = output_origin_flat_index + ti * output_dims[2] + tj; - int output_increment = WriteRowPerPass * output_dims[2]; + int64_t output_index = output_origin_flat_index + + static_cast(ti) * output_dims[2] + tj; + int64_t output_increment = + static_cast(WriteRowPerPass) * output_dims[2]; if (full_tile) { #pragma unroll diff --git a/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc b/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc new file mode 100644 index 00000000000000..b84d73b9bf5b93 --- /dev/null +++ b/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc @@ -0,0 +1,37 @@ +/* Copyright 2026 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. +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 "tensorflow/core/kernels/conv_2d_gpu.h" + +#include + +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace functor { +namespace { + +TEST(TensorIndexToFlatTest, HandlesFlatIndexPastInt32Limit) { + constexpr int kDimension = 46341; + const Dimension<3> dims(1, kDimension, kDimension); + const Index<3> index(0, kDimension - 1, kDimension - 1); + + EXPECT_EQ(TensorIndexToFlat(index, dims), + static_cast(kDimension) * kDimension - 1); +} + +} // namespace +} // namespace functor +} // namespace tensorflow From bdb897354cb8dec00b736b0e81f33f6aced0b5a4 Mon Sep 17 00:00:00 2001 From: kokol16 Date: Wed, 9 Sep 2026 12:53:07 -0400 Subject: [PATCH 06/25] Fix conv_2d GPU test build Host builds compiled the GPU test without a CUDA or ROCm guard, and the test included conv_2d_gpu.h without listing it as a source. TensorFlow's build checks rejected both cases. I added the GPU guard and listed conv_2d_gpu.h in the test sources. --- tensorflow/core/kernels/BUILD | 5 ++++- tensorflow/core/kernels/conv_2d_gpu_test.cu.cc | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index c33068afa5a53f..b397b716d9329b 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -338,7 +338,10 @@ tf_kernel_library( tf_cuda_only_cc_test( name = "conv_2d_gpu_test", size = "small", - srcs = ["conv_2d_gpu_test.cu.cc"], + srcs = [ + "conv_2d_gpu.h", + "conv_2d_gpu_test.cu.cc", + ], features = if_cuda(["-layering_check"]), deps = [ ":conv_2d", diff --git a/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc b/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc index b84d73b9bf5b93..112f812a0a51b6 100644 --- a/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc +++ b/tensorflow/core/kernels/conv_2d_gpu_test.cu.cc @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + #include "tensorflow/core/kernels/conv_2d_gpu.h" #include @@ -35,3 +37,5 @@ TEST(TensorIndexToFlatTest, HandlesFlatIndexPastInt32Limit) { } // namespace } // namespace functor } // namespace tensorflow + +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM From 88d6390db4fdcc6d20d2ae381f24841e97b554f0 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 29 Sep 2026 04:35:04 -0700 Subject: [PATCH 07/25] Squashing commits after sync tool breakage Reverts changelist 989948905 PiperOrigin-RevId: 990219345 --- .../lite/tests/optimize_batch_matmul.mlir | 64 ++ .../transforms/optimize_batch_matmul_pass.cc | 41 +- tensorflow/compiler/mlir/tfr/ir/tfr_ops.td | 1 + .../mlir/tfrt/ir/mlrt/mlrt_dialect.td | 1 + .../mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td | 1 + .../mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td | 1 + .../mlir/tfrt/ir/tfrt_fallback_async.td | 1 + .../mlir/tfrt/ir/tfrt_fallback_sync.td | 1 + .../tools/kernel_gen/ir/tf_framework_ops.td | 1 + tensorflow/core/ir/dialect.td | 1 + tensorflow/dtensor/python/tests/test_util.py | 4 +- tensorflow/lite/core/c/BUILD | 7 +- tensorflow/lite/delegates/xnnpack/BUILD | 41 ++ .../lite/delegates/xnnpack/moe_block_scale.h | 120 ++++ .../delegates/xnnpack/moe_block_scale_test.cc | 244 ++++++++ .../delegates/xnnpack/moe_delegate_kernel.cc | 78 +-- .../xnnpack/moe_delegate_kernel_test.cc | 217 +++++++ .../lite/kernels/internal/reference/select.h | 10 +- tensorflow/lite/python/analyzer_wrapper/BUILD | 1 - .../reorder_reshape_transpose.cc | 1 - tensorflow/lite/tools/strip_buffers/BUILD | 3 - .../reconstitute_buffers_into_fb.cc | 1 - .../strip_buffers/strip_buffers_from_fb.cc | 1 - .../lite/tools/strip_buffers/stripping_lib.cc | 1 - .../service/cross_trainer_cache_test.py | 4 +- .../control_flow/control_flow_ops_py_test.py | 72 +++ .../kernel_tests/math_ops/bincount_op_test.py | 24 + tensorflow/python/ops/BUILD | 3 + tensorflow/python/ops/control_flow_assert.py | 54 +- tensorflow/workspace2.bzl | 2 + .../xla/.github/workflows/ci_multi_device.yml | 4 +- third_party/xla/AGENTS.md | 11 + third_party/xla/docs/custom_call.md | 2 +- third_party/xla/opensource_only.files | 1 + .../cuda_tile/patches/llvm_cl_989877215.patch | 11 + .../xla/third_party/cuda_tile/patches/series | 1 + .../xla/third_party/extensions/llvm.bzl | 2 + .../xla/third_party/llvm/toolchains.patch | 19 +- .../xla/third_party/triton/common/series.bzl | 1 + .../useStrictPropertiesInAssemblyFormat.patch | 114 ++++ third_party/xla/third_party/xz/BUILD.bazel | 0 .../xla/third_party/xz/config/BUILD.bazel | 1 + .../third_party/xz/config/config_generic.h | 557 ++++++++++++++++++ .../xla/third_party/xz/patches/BUILD.bazel | 1 + .../xz/patches/remove_have_config.diff | 14 + third_party/xla/third_party/xz/workspace.bzl | 37 ++ third_party/xla/third_party/xz/xz.BUILD.bazel | 172 ++++++ third_party/xla/workspace2.bzl | 2 + third_party/xla/xla/AGENTS.md | 56 +- .../gpu/autotuner/block_level_emitter.cc | 15 +- .../backends/gpu/autotuner/fission_backend.cc | 3 +- .../xla/xla/backends/gpu/autotuner/miopen.cc | 36 +- .../xla/xla/backends/gpu/autotuner/miopen.h | 15 + .../xla/backends/gpu/autotuner/miopen_test.cc | 5 +- .../gpu/codegen/triton/collective_emitter.cc | 133 +++-- .../gpu/codegen/triton/collective_emitter.h | 17 +- .../gpu/codegen/triton/ir/tests/invalid.mlir | 4 +- .../codegen/triton/ir/triton_xla_dialect.td | 1 + .../gpu/codegen/triton/ir/triton_xla_ops.td | 6 +- .../backends/gpu/codegen/triton/tests/BUILD | 1 - .../triton/tests/collectives/all_gather.hlo | 25 +- .../tests/collectives/all_gather_one_shot.hlo | 58 -- .../triton/tests/collectives/all_reduce.hlo | 4 +- .../fusion_emitter_shared_dialect_test.cc | 4 +- .../transforms/stablehlo_lower_to_triton.cc | 15 +- .../tests/triton_xla_block_barrier.mlir | 4 +- .../tests/triton_xla_fold_transpose.mlir | 56 ++ .../tests/triton_xla_remote_access.mlir | 8 +- .../tests/triton_xla_squeeze_dims.mlir | 46 +- .../triton_xla_fold_transpose_pass.cc | 129 +++- .../gpu/codegen/triton/xtile_compiler.cc | 8 - .../xla/xla/backends/gpu/runtime/BUILD | 58 ++ .../xla/backends/gpu/runtime/all_gather.cc | 1 - .../xla/backends/gpu/runtime/all_reduce.cc | 1 - .../runtime/collective_kernel_thunk_test.cc | 1 - .../backends/gpu/runtime/collective_params.h | 3 - .../runtime/command_buffer_conversion_pass.cc | 3 +- .../xla/backends/gpu/runtime/cow_storage.cc | 78 +++ .../xla/backends/gpu/runtime/cow_storage.h | 80 +++ .../backends/gpu/runtime/cow_storage_test.cc | 304 ++++++++++ .../xla/backends/gpu/runtime/device_slot.h | 59 ++ .../backends/gpu/runtime/vector_storage.cc | 29 + .../xla/backends/gpu/runtime/vector_storage.h | 49 ++ .../gpu/runtime/vector_storage_test.cc | 138 +++++ third_party/xla/xla/backends/gpu/tests/BUILD | 4 + .../backends/gpu/tests/all_gather_e2e_test.cc | 201 ++++++- .../backends/gpu/tests/all_reduce_e2e_test.cc | 200 ++++++- .../gpu/tests/collective_ops_e2e_test_base.cc | 26 +- .../gpu/tests/swap_conv_operands_test.cc | 8 + .../xla/xla/backends/gpu/transforms/BUILD | 2 +- .../transforms/fusion_block_level_rewriter.cc | 7 +- .../gpu/transforms/layout_assignment.cc | 39 +- .../gpu/transforms/layout_assignment_a100.hlo | 15 +- .../gpu/transforms/layout_assignment_h100.hlo | 15 +- .../gpu/transforms/layout_assignment_v100.hlo | 15 +- .../gpu/transforms/priority_fusion.cc | 3 +- .../gpu/transforms/priority_fusion_test.cc | 2 +- .../gpu/transforms/softmax_rewriter_triton.cc | 7 +- .../xla/xla/backends/profiler/cpu/BUILD | 4 +- .../backends/profiler/cpu/python_tracer.cc | 40 +- .../xla/backends/profiler/cpu/python_tracer.h | 7 +- .../tiling/experimental/tile_propagation.cc | 63 ++ .../experimental/tile_propagation_test.cc | 155 +++++ .../tiling/experimental/tiled_hlo_test.cc | 41 ++ .../codegen/experimental_fusion_emitter.cc | 50 +- .../ir/transforms/verify_legal_xtile_ops.cc | 24 +- .../xla/xla/codegen/xtile/ir/xtile_dialect.td | 1 + third_party/xla/xla/debug_options_flags.cc | 3 +- .../xla/xla/debug_options_parsers_test.cc | 31 +- .../xla/hlo/analysis/hlo_dataflow_analysis.cc | 131 +++- .../xla/hlo/analysis/hlo_dataflow_analysis.h | 7 +- .../analysis/hlo_dataflow_analysis_test.cc | 104 ++++ third_party/xla/xla/hlo/ir/hlo_instructions.h | 70 ++- .../xla/xla/hlo/transforms/simplifiers/BUILD | 2 + .../simplifiers/sort_simplifier_test.cc | 92 ++- .../xla/mlir_hlo/mhlo/IR/hlo_ops_common.td | 1 + third_party/xla/xla/mosaic/dialect/tpu/tpu.td | 1 + .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc | 37 ++ .../pjrt/se/pjrt_stream_executor_client.cc | 25 +- third_party/xla/xla/pjrt/se/se_raw_buffer.cc | 32 + third_party/xla/xla/python/_profiler.pyi | 7 +- .../xla/xla/python/ifrt/ir/ifrt_dialect.td | 1 + .../xla/xla/python/ifrt/ir/vifrt_dialect.td | 1 + third_party/xla/xla/python/profiler.cc | 50 +- .../xla/python/profiler/profile_data_lib.cc | 16 +- .../xla/python/profiler/profile_data_lib.h | 18 +- third_party/xla/xla/service/BUILD | 8 +- ...nc_collective_custom_call_rewriter_test.cc | 23 +- third_party/xla/xla/service/copy_insertion.cc | 78 ++- third_party/xla/xla/service/copy_insertion.h | 10 +- .../xla/xla/service/copy_insertion_test.cc | 131 ++++ third_party/xla/xla/service/gpu/BUILD | 30 + .../gpu/gpu_executable_buffer_allocator.cc | 51 +- .../gpu/gpu_executable_buffer_allocator.h | 4 +- .../gpu_executable_buffer_allocator_test.cc | 241 ++++++++ third_party/xla/xla/service/gpu/model/BUILD | 7 + .../model/combined_gpu_performance_model.cc | 24 +- .../model/combined_gpu_performance_model.h | 26 +- .../model/gpu_indexing_performance_model.cc | 466 +++++++++------ .../model/gpu_indexing_performance_model.h | 33 +- .../gpu_indexing_performance_model_test.cc | 82 ++- .../xla/xla/service/layout_assignment.cc | 298 +++++----- .../xla/xla/service/layout_assignment.h | 67 +-- .../xla/xla/service/layout_assignment_test.cc | 4 +- third_party/xla/xla/stream_executor/BUILD | 2 +- .../xla/xla/stream_executor/rocm/rocm_dnn.cc | 37 +- .../xla/xla/stream_executor/rocm/rocm_dnn.h | 13 + .../functional_hlo_runner.cc | 127 ++-- third_party/xla/xla/xla.proto | 1 + 149 files changed, 5683 insertions(+), 1042 deletions(-) create mode 100644 tensorflow/lite/delegates/xnnpack/moe_block_scale.h create mode 100644 tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc create mode 100644 tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc create mode 100644 third_party/xla/AGENTS.md create mode 100644 third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch create mode 100644 third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch create mode 100644 third_party/xla/third_party/xz/BUILD.bazel create mode 100644 third_party/xla/third_party/xz/config/BUILD.bazel create mode 100644 third_party/xla/third_party/xz/config/config_generic.h create mode 100644 third_party/xla/third_party/xz/patches/BUILD.bazel create mode 100644 third_party/xla/third_party/xz/patches/remove_have_config.diff create mode 100644 third_party/xla/third_party/xz/workspace.bzl create mode 100644 third_party/xla/third_party/xz/xz.BUILD.bazel delete mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo create mode 100644 third_party/xla/xla/backends/gpu/runtime/cow_storage.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/cow_storage.h create mode 100644 third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/device_slot.h create mode 100644 third_party/xla/xla/backends/gpu/runtime/vector_storage.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/vector_storage.h create mode 100644 third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc create mode 100644 third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc diff --git a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir index 8c15335815bf4b..8b1829faf66acc 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir @@ -342,4 +342,68 @@ func.func @Batchmatmul2FullyconnectedReshapeTransposeChainConstY(%arg0: tensor<1 // CHECK-NOT: "tfl.batch_matmul" %2 = "tfl.batch_matmul"(%arg0, %1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> func.return %2 : tensor<16x128xf32> +} + +// A bf16 constant widened to f32 is still foldable, since flatbuffer export +// materializes the cast. The transpose goes below the cast so it lands +// directly on the constant, where it can be folded. +// CHECK-LABEL: Batchmatmul2FullyconnectedBf16CastConstY +func.func @Batchmatmul2FullyconnectedBf16CastConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16> + %0 = "tfl.cast"(%weight) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32> + // CHECK: %[[W:.*]] = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16> + // CHECK: %[[T:.*]] = "tfl.transpose"(%[[W]], %{{.*}}) : (tensor<1024x128xbf16>, tensor<2xi32>) -> tensor<128x1024xbf16> + // CHECK: %[[C:.*]] = "tfl.cast"(%[[T]]) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32> + // CHECK: "tfl.fully_connected"(%arg0, %[[C]], %{{.*}}) + // CHECK-NOT: "tfl.batch_matmul" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// With adj_y the rhs is already in filter layout, so the cast is kept as is. +// CHECK-LABEL: Batchmatmul2FullyconnectedBf16CastConstYAdjY +func.func @Batchmatmul2FullyconnectedBf16CastConstYAdjY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %weight = arith.constant dense_resource<__elided__> : tensor<128x1024xbf16> + %0 = "tfl.cast"(%weight) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32> + // CHECK: %[[W:.*]] = arith.constant dense_resource<__elided__> : tensor<128x1024xbf16> + // CHECK: %[[C:.*]] = "tfl.cast"(%[[W]]) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32> + // CHECK-NOT: "tfl.transpose" + // CHECK: "tfl.fully_connected"(%arg0, %[[C]], %{{.*}}) + // CHECK-NOT: "tfl.batch_matmul" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = true, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<128x1024xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulBf16CastConstant +func.func @NotFuseTransposeFCRhsToBatchMatmulBf16CastConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> { + %cst = arith.constant dense<[1, 0]> : tensor<2xi32> + %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16> + %0 = "tfl.cast"(%weight) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32> + %1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32> + // CHECK: "tfl.fully_connected" + // CHECK-NOT: "tfl.batch_matmul" + %2 = "tfl.fully_connected"(%arg0, %1, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32> + func.return %2 : tensor<16x128xf32> +} + +// A widening cast of a non-constant is not foldable. +// CHECK-LABEL: NotBatchmatmul2FullyconnectedBf16CastNonConstY +func.func @NotBatchmatmul2FullyconnectedBf16CastNonConstY(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xbf16>) -> tensor<16x128xf32> { + %0 = "tfl.cast"(%arg1) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32> + // CHECK: "tfl.batch_matmul" + // CHECK-NOT: "tfl.fully_connected" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> +} + +// Only float widening is looked through; flatbuffer export does not +// materialize other casts. +// CHECK-LABEL: NotBatchmatmul2FullyconnectedI8CastConstY +func.func @NotBatchmatmul2FullyconnectedI8CastConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> { + %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xi8> + %0 = "tfl.cast"(%weight) : (tensor<1024x128xi8>) -> tensor<1024x128xf32> + // CHECK: "tfl.batch_matmul" + // CHECK-NOT: "tfl.fully_connected" + %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32> + func.return %1 : tensor<16x128xf32> } \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index c8e4ce5c6077d9..35bdffefc50ff2 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -44,9 +44,26 @@ namespace mlir { namespace TFL { namespace { +// Returns `op` as a bf16/f16 -> f32 cast, or null if it is anything else. +// +// Weights kept in a narrow float type reach the matmul behind such a cast, +// since the kernels only run in f32. Flatbuffer export materializes this cast +// when its input is a constant, so the cast does not stop the value from being +// a constant filter. +CastOp AsFloatWideningCast(mlir::Operation* op) { + auto cast_op = mlir::dyn_cast_or_null(op); + if (!cast_op) return nullptr; + mlir::Type in_elem = getElementTypeOrSelf(cast_op.getInput().getType()); + mlir::Type out_elem = getElementTypeOrSelf(cast_op.getType()); + if (!(in_elem.isBF16() || in_elem.isF16()) || !out_elem.isF32()) { + return nullptr; + } + return cast_op; +} + // Checks whether `value` is, or can be folded into, a constant. This covers a // plain constant, the dequantize wrappers a quantized constant weight hides -// behind, and chains of ops that only rearrange one of those. +// behind, and chains of ops that only rearrange or widen one of those. bool IsFromFoldableChain(mlir::Value value) { mlir::Operation* defining_op = value.getDefiningOp(); @@ -66,6 +83,8 @@ bool IsFromFoldableChain(mlir::Value value) { defining_op = split_op.getValue().getDefiningOp(); } else if (auto transpose_op = mlir::dyn_cast(defining_op)) { defining_op = transpose_op.getInput().getDefiningOp(); + } else if (auto cast_op = AsFloatWideningCast(defining_op)) { + defining_op = cast_op.getInput().getDefiningOp(); } else { // Stop if the op does not preserve constant nature. break; @@ -193,8 +212,24 @@ struct ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs // The rhs need to be transposed if adj_y == false AND this matmul will be // legalized to tfl.fully_connected - Value output_rhs = - !bmm_op.getAdjY() ? create_z_x_transpose_op(input_rhs) : input_rhs; + Value output_rhs = input_rhs; + if (!bmm_op.getAdjY()) { + // Transpose below a widening cast rather than above it. Only a + // transpose directly on a constant gets folded away, and the cast is + // then left directly on the constant too, which flatbuffer export + // materializes. + if (CastOp cast_op = AsFloatWideningCast(input_rhs.getDefiningOp())) { + Value transposed = create_z_x_transpose_op(cast_op.getInput()); + auto transposed_type = + mlir::cast(transposed.getType()); + output_rhs = TFL::CastOp::create( + rewriter, cast_op.getLoc(), + transposed_type.clone(getElementTypeOrSelf(cast_op.getType())), + transposed); + } else { + output_rhs = create_z_x_transpose_op(input_rhs); + } + } Type output_type = bmm_op.getResult().getType(); auto no_input = diff --git a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td index 2cad28974ab8ec..be5401d5fabc79 100644 --- a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td +++ b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td @@ -41,6 +41,7 @@ def TFR_Dialect : Dialect { }]; let cppNamespace = "::mlir::TFR"; + let useStrictPropertiesInAssemblyFormat = 0; } //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td b/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td index 13409c3ece1f3d..2650a412198272 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td +++ b/tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td @@ -26,6 +26,7 @@ def Mlrt_Dialect : Dialect { }]; let cppNamespace = "::mlrt::compiler"; + let useStrictPropertiesInAssemblyFormat = 0; } def MlrtFutureType : DialectType traits = []> : diff --git a/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td b/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td index a207b83c7e5207..f545786faead7a 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td +++ b/tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td @@ -28,6 +28,7 @@ def TensorflowMlrtTpu_Dialect : Dialect { }]; let cppNamespace = "::tensorflow::tf_mlrt_tpu"; + let useStrictPropertiesInAssemblyFormat = 0; } class TensorflowMlrtTpu_Op traits = []> : diff --git a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td index 3209f2f253cc2a..7d73e733ffcb97 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td +++ b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td @@ -33,6 +33,7 @@ def FallbackAsync_Dialect : Dialect { }]; let cppNamespace = "::tfrt::fallback_async"; + let useStrictPropertiesInAssemblyFormat = 0; } class FallbackAsync_Op traits = []> : diff --git a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td index bba8a0219217c6..4575efae94871f 100644 --- a/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td +++ b/tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td @@ -32,6 +32,7 @@ def FallbackSync_Dialect : Dialect { }]; let cppNamespace = "::tfrt::fallback_sync"; + let useStrictPropertiesInAssemblyFormat = 0; } class FallbackSync_Op traits = []> : diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td b/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td index 64f782d02346e8..b5e11c53ddd97e 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td +++ b/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td @@ -40,6 +40,7 @@ def TFFramework_Dialect : Dialect { }]; let useDefaultTypePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; } def TFFramework_OpKernelContextType : DialectType +#include +#include + +// Scale layout and dequantization helpers shared by the custom "moe" XNNPACK +// delegate kernel. They live in their own header so they can be unit tested +// without standing up a delegate. +// +// MoE expert weights are laid out as `[output_channels, num_experts, 1, +// input_channels]`, i.e. one row per (output channel, expert) pair. The +// matching scale tensor is `[output_channels, num_experts, 1, blocks_per_row]`. +// `blocks_per_row == 1` is per-output-channel quantization; larger values split +// the input axis into that many equally sized blocks, each with its own scale. +namespace tflite { +namespace xnnpack { + +// How a single row's scales map onto its input channels. +struct BlockScaleLayout { + // Number of scales per row. 1 means per-output-channel quantization. + size_t groups_per_row; + // Number of input channels covered by one scale. + size_t group_size; +}; + +// Derives the layout from the total number of scale elements. `scale_elements` +// is expected to be an exact multiple of `num_rows`; callers validate that +// before reaching here, and a ragged value degrades to treating the row as +// per-channel rather than reading out of bounds. +inline BlockScaleLayout ResolveBlockScaleLayout(size_t scale_elements, + size_t num_rows, + size_t input_channels) { + const size_t groups_per_row = + (num_rows > 0) ? (scale_elements / num_rows) : 0; + const size_t group_size = + (groups_per_row > 0 && groups_per_row <= input_channels) + ? (input_channels / groups_per_row) + : 1; + return {groups_per_row, group_size}; +} + +// Index of the scale covering input channel `in`. Clamped to the last block so +// an input axis that does not divide evenly cannot read past the row. +inline size_t BlockScaleIndex(const BlockScaleLayout& layout, size_t in) { + if (layout.groups_per_row <= 1 || layout.group_size == 0) { + return 0; + } + return std::min(in / layout.group_size, layout.groups_per_row - 1); +} + +// Dequantizes the rows belonging to `expert` into `dst`, which is a dense +// `[output_channels, input_channels]` buffer. +inline void CopyAndDequantizeExpertWeightRowsInt8( + const int8_t* weight_i8, const float* scale, size_t scale_elements, + size_t num_experts, size_t expert, size_t output_channels, + size_t input_channels, float* dst) { + const size_t num_rows = output_channels * num_experts; + const BlockScaleLayout layout = + ResolveBlockScaleLayout(scale_elements, num_rows, input_channels); + + for (size_t out = 0; out < output_channels; ++out) { + const size_t row_idx = out * num_experts + expert; + const int8_t* src_row = weight_i8 + row_idx * input_channels; + float* dst_row = dst + out * input_channels; + const float* row_scales = scale + row_idx * layout.groups_per_row; + for (size_t in = 0; in < input_channels; ++in) { + const float scale_val = row_scales[BlockScaleIndex(layout, in)]; + dst_row[in] = static_cast(src_row[in]) * scale_val; + } + } +} + +// As above, for int4 weights packed two nibbles per byte, low nibble first. +inline void CopyAndDequantizeExpertWeightRowsInt4( + const int8_t* weight_i4_packed, const float* scale, size_t scale_elements, + size_t num_experts, size_t expert, size_t output_channels, + size_t input_channels, float* dst) { + const size_t num_rows = output_channels * num_experts; + const BlockScaleLayout layout = + ResolveBlockScaleLayout(scale_elements, num_rows, input_channels); + + for (size_t out = 0; out < output_channels; ++out) { + const size_t row_idx = out * num_experts + expert; + const int8_t* src_row_packed = + weight_i4_packed + (row_idx * input_channels) / 2; + float* dst_row = dst + out * input_channels; + const float* row_scales = scale + row_idx * layout.groups_per_row; + + for (size_t in = 0; in < input_channels; ++in) { + const size_t byte_idx = in / 2; + const int8_t byte_val = src_row_packed[byte_idx]; + const int8_t nibble = (in % 2 == 0) + ? static_cast(byte_val << 4) >> 4 + : static_cast(byte_val >> 4); + const float scale_val = row_scales[BlockScaleIndex(layout, in)]; + dst_row[in] = static_cast(nibble) * scale_val; + } + } +} + +} // namespace xnnpack +} // namespace tflite + +#endif // TENSORFLOW_LITE_DELEGATES_XNNPACK_MOE_BLOCK_SCALE_H_ diff --git a/tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc b/tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc new file mode 100644 index 00000000000000..45ca9dc0b5b9a0 --- /dev/null +++ b/tensorflow/lite/delegates/xnnpack/moe_block_scale_test.cc @@ -0,0 +1,244 @@ +/* Copyright 2026 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. +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 "tensorflow/lite/delegates/xnnpack/moe_block_scale.h" + +#include +#include +#include + +#include +#include + +namespace tflite { +namespace xnnpack { +namespace { + +using ::testing::ElementsAre; + +TEST(ResolveBlockScaleLayoutTest, OneScalePerRowIsPerChannel) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/8, /*num_rows=*/8, /*input_channels=*/16); + EXPECT_EQ(layout.groups_per_row, 1); + EXPECT_EQ(layout.group_size, 16); +} + +TEST(ResolveBlockScaleLayoutTest, SplitsInputAxisIntoEqualBlocks) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/32, /*num_rows=*/8, /*input_channels=*/64); + EXPECT_EQ(layout.groups_per_row, 4); + EXPECT_EQ(layout.group_size, 16); +} + +// A scale per input channel is the finest blocking the layout can express. +TEST(ResolveBlockScaleLayoutTest, HandlesOneScalePerInputChannel) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/64, /*num_rows=*/8, /*input_channels=*/8); + EXPECT_EQ(layout.groups_per_row, 8); + EXPECT_EQ(layout.group_size, 1); +} + +// More scales than input channels cannot describe a real blocking. The parser +// rejects such tensors, but the layout still has to stay in bounds rather than +// compute a zero group size and divide by it. +TEST(ResolveBlockScaleLayoutTest, ClampsGroupSizeWhenBlocksExceedChannels) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/64, /*num_rows=*/8, /*input_channels=*/4); + EXPECT_EQ(layout.groups_per_row, 8); + EXPECT_EQ(layout.group_size, 1); +} + +TEST(ResolveBlockScaleLayoutTest, ToleratesZeroRows) { + const BlockScaleLayout layout = ResolveBlockScaleLayout( + /*scale_elements=*/0, /*num_rows=*/0, /*input_channels=*/16); + EXPECT_EQ(layout.groups_per_row, 0); + EXPECT_EQ(layout.group_size, 1); +} + +TEST(BlockScaleIndexTest, PerChannelLayoutAlwaysUsesTheFirstScale) { + const BlockScaleLayout layout = {/*groups_per_row=*/1, /*group_size=*/16}; + EXPECT_EQ(BlockScaleIndex(layout, 0), 0); + EXPECT_EQ(BlockScaleIndex(layout, 15), 0); +} + +TEST(BlockScaleIndexTest, AdvancesAtBlockBoundaries) { + const BlockScaleLayout layout = {/*groups_per_row=*/4, /*group_size=*/4}; + std::vector indices; + indices.reserve(16); + for (size_t in = 0; in < 16; ++in) { + indices.push_back(BlockScaleIndex(layout, in)); + } + EXPECT_THAT(indices, + ElementsAre(0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3)); +} + +// Reading past the last block would walk into the next row's scales, so the +// index saturates instead. +TEST(BlockScaleIndexTest, ClampsToTheLastBlock) { + const BlockScaleLayout layout = {/*groups_per_row=*/4, /*group_size=*/4}; + EXPECT_EQ(BlockScaleIndex(layout, 16), 3); + EXPECT_EQ(BlockScaleIndex(layout, 1000), 3); +} + +TEST(BlockScaleIndexTest, ToleratesZeroGroupSize) { + const BlockScaleLayout layout = {/*groups_per_row=*/4, /*group_size=*/0}; + EXPECT_EQ(BlockScaleIndex(layout, 7), 0); +} + +// Two experts, two output channels, four input channels. Rows are interleaved +// as [out, expert], which is what makes the expert stride `num_experts`. +class DequantizeInt8Test : public ::testing::Test { + protected: + static constexpr size_t kNumExperts = 2; + static constexpr size_t kOutputChannels = 2; + static constexpr size_t kInputChannels = 4; + + // Row r holds the values {1, 2, 3, 4} offset by 10 * r, so a mis-indexed row + // is obvious in the output. + std::vector Weights() const { + std::vector weights; + for (size_t row = 0; row < kOutputChannels * kNumExperts; ++row) { + for (size_t in = 0; in < kInputChannels; ++in) { + weights.push_back(static_cast(10 * row + in + 1)); + } + } + return weights; + } +}; + +TEST_F(DequantizeInt8Test, PerChannelScalesApplyToTheWholeRow) { + const std::vector weights = Weights(); + // One scale per row; rows are [out0/e0, out0/e1, out1/e0, out1/e1]. + const std::vector scales = {1.0f, 100.0f, 2.0f, 200.0f}; + std::vector dst(kOutputChannels * kInputChannels, -1.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + // Expert 0 owns rows 0 and 2, scaled by 1.0 and 2.0. + EXPECT_THAT(dst, ElementsAre(1.0f, 2.0f, 3.0f, 4.0f, // row 0 * 1.0 + 42.0f, 44.0f, 46.0f, 48.0f)); // row 2 * 2.0 +} + +TEST_F(DequantizeInt8Test, BlockwiseScalesChangeHalfwayThroughTheRow) { + const std::vector weights = Weights(); + // Two scales per row over four input channels, so a block is two channels + // wide. Rows are [out0/e0, out0/e1, out1/e0, out1/e1]. + const std::vector scales = {1.0f, 10.0f, // row 0 + -1.0f, -10.0f, // row 1 (other expert) + 2.0f, 20.0f, // row 2 + -2.0f, -20.0f}; // row 3 (other expert) + std::vector dst(kOutputChannels * kInputChannels, -1.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre( + // row 0 = {1,2,3,4}, scales {1, 1, 10, 10} + 1.0f, 2.0f, 30.0f, 40.0f, + // row 2 = {21,22,23,24}, scales {2, 2, 20, 20} + 42.0f, 44.0f, 460.0f, 480.0f)); +} + +// The second expert must read its own interleaved rows and its own scales. +TEST_F(DequantizeInt8Test, BlockwiseScalesAreSelectedPerExpert) { + const std::vector weights = Weights(); + const std::vector scales = {1.0f, 10.0f, -1.0f, -10.0f, + 2.0f, 20.0f, -2.0f, -20.0f}; + std::vector dst(kOutputChannels * kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/1, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre( + // row 1 = {11,12,13,14}, scales {-1, -1, -10, -10} + -11.0f, -12.0f, -130.0f, -140.0f, + // row 3 = {31,32,33,34}, scales {-2, -2, -20, -20} + -62.0f, -64.0f, -660.0f, -680.0f)); +} + +// One scale per input channel is the limiting case of blocking. +TEST_F(DequantizeInt8Test, SupportsOneScalePerInputChannel) { + const std::vector weights = Weights(); + std::vector scales(kOutputChannels * kNumExperts * kInputChannels, + 0.0f); + // Only expert 0's rows (0 and 2) matter here. + for (size_t i = 0; i < kInputChannels; ++i) { + scales[0 * kInputChannels + i] = static_cast(i + 1); + scales[2 * kInputChannels + i] = static_cast(i + 1); + } + std::vector dst(kOutputChannels * kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt8( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre( + // row 0 = {1,2,3,4} times {1,2,3,4} + 1.0f, 4.0f, 9.0f, 16.0f, + // row 2 = {21,22,23,24} times {1,2,3,4} + 21.0f, 44.0f, 69.0f, 96.0f)); +} + +class DequantizeInt4Test : public ::testing::Test { + protected: + static constexpr size_t kNumExperts = 2; + static constexpr size_t kOutputChannels = 1; + static constexpr size_t kInputChannels = 4; + + // Two nibbles per byte, low nibble first. Row 0 holds {1, 2, 3, 4} and row 1 + // holds {-1, -2, -3, -4}. + std::vector PackedWeights() const { + return { + static_cast(0x21), + static_cast(0x43), // row 0 + static_cast(0xEF), + static_cast(0xCD), // row 1 + }; + } +}; + +TEST_F(DequantizeInt4Test, UnpacksNibblesWithPerChannelScale) { + const std::vector weights = PackedWeights(); + const std::vector scales = {2.0f, 3.0f}; + std::vector dst(kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt4( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/0, + kOutputChannels, kInputChannels, dst.data()); + + EXPECT_THAT(dst, ElementsAre(2.0f, 4.0f, 6.0f, 8.0f)); +} + +TEST_F(DequantizeInt4Test, AppliesBlockwiseScalesToNibbles) { + const std::vector weights = PackedWeights(); + // Two scales per row over four input channels. + const std::vector scales = {1.0f, 10.0f, // row 0 + 5.0f, 50.0f}; // row 1 + std::vector dst(kInputChannels, 0.0f); + + CopyAndDequantizeExpertWeightRowsInt4( + weights.data(), scales.data(), scales.size(), kNumExperts, /*expert=*/1, + kOutputChannels, kInputChannels, dst.data()); + + // Row 1 = {-1, -2, -3, -4} with scales {5, 5, 50, 50}. + EXPECT_THAT(dst, ElementsAre(-5.0f, -10.0f, -150.0f, -200.0f)); +} + +} // namespace +} // namespace xnnpack +} // namespace tflite diff --git a/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc index 27421f5dca40fb..bc06fd34069150 100644 --- a/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc +++ b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.cc @@ -34,6 +34,7 @@ limitations under the License. #include "tensorflow/lite/c/builtin_op_data.h" #include "tensorflow/lite/core/c/builtin_op_data.h" #include "tensorflow/lite/core/c/common.h" +#include "tensorflow/lite/delegates/xnnpack/moe_block_scale.h" #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/kernel_util.h" @@ -185,15 +186,11 @@ class MoeExpertsDelegateKernel::Impl { return false; } const int count = NumElements(t); - if (attr.weight_type == MoeExpertsAttributes::WeightType::kInt8) { - if (count != static_cast(num_rows)) { - TF_LITE_KERNEL_LOG( - context, - "%s node #%d %s element count %d does not match rows %zu", - kMoeCustomOp, node_index, name, count, num_rows); - return false; - } - } else if (attr.weight_type == MoeExpertsAttributes::WeightType::kInt4) { + // One scale per row is per-output-channel quantization; any exact + // multiple of the row count is blockwise, with that multiple being the + // number of blocks along the input axis. + if (attr.weight_type == MoeExpertsAttributes::WeightType::kInt8 || + attr.weight_type == MoeExpertsAttributes::WeightType::kInt4) { if (count <= 0 || (static_cast(count) % num_rows != 0)) { TF_LITE_KERNEL_LOG( context, @@ -646,55 +643,6 @@ class MoeExpertsDelegateKernel::Impl { } } - static void CopyAndDequantizeExpertWeightRowsInt8( - const int8_t* weight_i8, const float* scale, size_t num_experts, - size_t expert, size_t output_channels, size_t input_channels, - float* dst) { - for (size_t out = 0; out < output_channels; ++out) { - const size_t row_idx = out * num_experts + expert; - const float row_scale = scale[row_idx]; - const int8_t* src_row = weight_i8 + row_idx * input_channels; - float* dst_row = dst + out * input_channels; - for (size_t in = 0; in < input_channels; ++in) { - dst_row[in] = static_cast(src_row[in]) * row_scale; - } - } - } - - static void CopyAndDequantizeExpertWeightRowsInt4( - const int8_t* weight_i4_packed, const float* scale, size_t scale_elements, - size_t num_experts, size_t expert, size_t output_channels, - size_t input_channels, float* dst) { - const size_t num_rows = output_channels * num_experts; - const size_t groups_per_row = scale_elements / num_rows; - const size_t group_size = - (groups_per_row > 0 && groups_per_row <= input_channels) - ? (input_channels / groups_per_row) - : 1; - - for (size_t out = 0; out < output_channels; ++out) { - const size_t row_idx = out * num_experts + expert; - const int8_t* src_row_packed = - weight_i4_packed + (row_idx * input_channels) / 2; - float* dst_row = dst + out * input_channels; - const float* row_scales = scale + row_idx * groups_per_row; - - for (size_t in = 0; in < input_channels; ++in) { - const size_t byte_idx = in / 2; - const int8_t byte_val = src_row_packed[byte_idx]; - const int8_t nibble = (in % 2 == 0) - ? static_cast(byte_val << 4) >> 4 - : static_cast(byte_val >> 4); - const size_t group_idx = - (groups_per_row > 1 && group_size > 0) - ? std::min(in / group_size, groups_per_row - 1) - : 0; - const float scale_val = row_scales[group_idx]; - dst_row[in] = static_cast(nibble) * scale_val; - } - } - } - void CopyGateUpExpertWeight(const void* gate_weight, const float* gate_scale, size_t gate_scale_elements, const void* ff1_weight, const float* ff1_scale, @@ -715,11 +663,12 @@ class MoeExpertsDelegateKernel::Impl { dst + hidden_dim * model_dim); } else if (attr_.weight_type == MoeExpertsAttributes::WeightType::kInt8) { CopyAndDequantizeExpertWeightRowsInt8( - static_cast(gate_weight), gate_scale, num_experts, - expert, hidden_dim, model_dim, dst); + static_cast(gate_weight), gate_scale, + gate_scale_elements, num_experts, expert, hidden_dim, model_dim, dst); CopyAndDequantizeExpertWeightRowsInt8( - static_cast(ff1_weight), ff1_scale, num_experts, - expert, hidden_dim, model_dim, dst + hidden_dim * model_dim); + static_cast(ff1_weight), ff1_scale, ff1_scale_elements, + num_experts, expert, hidden_dim, model_dim, + dst + hidden_dim * model_dim); } else { CopyExpertWeightRows(static_cast(gate_weight), num_experts, expert, hidden_dim, model_dim, dst); @@ -741,8 +690,9 @@ class MoeExpertsDelegateKernel::Impl { kernel_buffer_.data()); } else if (attr_.weight_type == MoeExpertsAttributes::WeightType::kInt8) { CopyAndDequantizeExpertWeightRowsInt8( - static_cast(weight), scale, num_experts, expert, - output_channels, input_channels, kernel_buffer_.data()); + static_cast(weight), scale, scale_elements, + num_experts, expert, output_channels, input_channels, + kernel_buffer_.data()); } else { CopyExpertWeightRows(static_cast(weight), num_experts, expert, output_channels, input_channels, diff --git a/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc new file mode 100644 index 00000000000000..a2f4430833f0d9 --- /dev/null +++ b/tensorflow/lite/delegates/xnnpack/moe_delegate_kernel_test.cc @@ -0,0 +1,217 @@ +/* Copyright 2026 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. +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 "tensorflow/lite/delegates/xnnpack/moe_delegate_kernel.h" + +#include +#include +#include +#include +#include + +#include +#include "flatbuffers/flexbuffers.h" // from @flatbuffers +#include "tensorflow/lite/builtin_ops.h" +#include "tensorflow/lite/core/c/common.h" +#include "tensorflow/lite/interpreter.h" + +namespace tflite { +namespace xnnpack { +namespace { + +constexpr int kNumExperts = 2; +constexpr int kNumActiveExperts = 1; +constexpr int kModelDim = 8; +constexpr int kHiddenDim = 4; +constexpr int kNumTokens = 2; + +// gate/ff1 are [hidden_dim, num_experts, 1, model_dim]; linear is the +// transpose. One scale per row is per-output-channel quantization. +constexpr int kGateFf1Rows = kHiddenDim * kNumExperts; +constexpr int kLinearRows = kModelDim * kNumExperts; + +std::vector MoeCustomData(const std::string& weight_type, + const std::string& activation) { + flexbuffers::Builder fbb; + fbb.Map([&]() { + fbb.Int("num_experts", kNumExperts); + fbb.Int("num_active_experts", kNumActiveExperts); + fbb.Int("model_dim", kModelDim); + fbb.Int("hidden_dim", kHiddenDim); + fbb.String("weight_type", weight_type); + fbb.String("activation", activation); + fbb.Bool("renormalized_top_weights", true); + }); + fbb.Finish(); + return fbb.GetBuffer(); +} + +// Builds a single-node interpreter holding one custom "moe" op so that +// MoeExpertsDelegateKernel::IsSupported can be called against a real +// TfLiteContext. The op is never invoked, so the tensor contents are +// irrelevant; only types, shapes and constness matter. +class MoeNodeBuilder { + public: + MoeNodeBuilder() : interpreter_(std::make_unique()) { + registration_ = {nullptr, nullptr, nullptr, nullptr}; + registration_.builtin_code = kTfLiteBuiltinCustom; + registration_.custom_name = "moe"; + registration_.version = 1; + } + + void AddInput(TfLiteType type, const std::vector& shape) { + interpreter_->AddTensors(1); + interpreter_->SetTensorParametersReadWrite( + next_tensor_, type, "", shape, {kTfLiteNoQuantization, nullptr}); + inputs_.push_back(next_tensor_++); + } + + void AddConstInput(TfLiteType type, const std::vector& shape, + size_t num_bytes) { + buffers_.push_back(std::vector(num_bytes, 0)); + interpreter_->AddTensors(1); + interpreter_->SetTensorParametersReadOnly( + next_tensor_, type, "", shape, {kTfLiteNoQuantization, nullptr}, + buffers_.back().data(), buffers_.back().size()); + inputs_.push_back(next_tensor_++); + } + + void AddOutput(TfLiteType type, const std::vector& shape) { + interpreter_->AddTensors(1); + interpreter_->SetTensorParametersReadWrite( + next_tensor_, type, "", shape, {kTfLiteNoQuantization, nullptr}); + outputs_.push_back(next_tensor_++); + } + + void SetCustomData(const std::vector& data) { + custom_data_.assign(data.begin(), data.end()); + } + + void Finish() { + interpreter_->SetInputs(inputs_); + interpreter_->SetOutputs(outputs_); + interpreter_->AddNodeWithParameters(inputs_, outputs_, custom_data_.data(), + custom_data_.size(), nullptr, + ®istration_); + } + + TfLiteStatus IsSupported() { + const auto* pair = interpreter_->node_and_registration(0); + return MoeExpertsDelegateKernel::IsSupported( + interpreter_->primary_subgraph().context(), &pair->first, &pair->second, + /*node_index=*/0); + } + + private: + std::unique_ptr interpreter_; + TfLiteRegistration registration_; + std::vector inputs_; + std::vector outputs_; + std::vector custom_data_; + // Backing storage for the constant tensors, which are referenced rather than + // copied by the interpreter. + std::vector> buffers_; + int next_tensor_ = 0; +}; + +// Assembles the ten inputs of an int8 moe op. `*_scales` is the total element +// count of each scale tensor: `rows` is per-output-channel and `rows * n` is +// blockwise with n blocks along the input axis. +void BuildInt8Moe(MoeNodeBuilder& builder, int gate_scales, int ff1_scales, + int linear_scales, + const std::string& activation = "gelu_tanh") { + builder.AddInput(kTfLiteFloat32, {1, 1, kNumTokens, kModelDim}); + builder.AddInput(kTfLiteFloat32, {1, 1, kNumTokens, kNumActiveExperts}); + builder.AddInput(kTfLiteInt32, {1, 1, kNumTokens, kNumActiveExperts}); + + builder.AddConstInput(kTfLiteInt8, {kHiddenDim, kNumExperts, 1, kModelDim}, + kGateFf1Rows * kModelDim * sizeof(int8_t)); + builder.AddConstInput(kTfLiteFloat32, {gate_scales}, + gate_scales * sizeof(float)); + builder.AddConstInput(kTfLiteInt8, {kHiddenDim, kNumExperts, 1, kModelDim}, + kGateFf1Rows * kModelDim * sizeof(int8_t)); + builder.AddConstInput(kTfLiteFloat32, {ff1_scales}, + ff1_scales * sizeof(float)); + builder.AddConstInput(kTfLiteInt8, {kModelDim, kNumExperts, 1, kHiddenDim}, + kLinearRows * kHiddenDim * sizeof(int8_t)); + builder.AddConstInput(kTfLiteFloat32, {linear_scales}, + linear_scales * sizeof(float)); + builder.AddConstInput(kTfLiteFloat32, {kNumExperts}, + kNumExperts * sizeof(float)); + + builder.AddOutput(kTfLiteFloat32, {1, 1, kNumTokens, kModelDim}); + builder.SetCustomData(MoeCustomData("int8", activation)); + builder.Finish(); +} + +TEST(MoeDelegateKernelIsSupportedTest, AcceptsPerChannelScales) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows, kGateFf1Rows, kLinearRows); + EXPECT_EQ(builder.IsSupported(), kTfLiteOk); +} + +// Four scales per row, i.e. the input axis split into four blocks. This is the +// case the delegate could not previously represent. +TEST(MoeDelegateKernelIsSupportedTest, AcceptsBlockwiseScales) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows * 4, kGateFf1Rows * 4, kLinearRows * 4); + EXPECT_EQ(builder.IsSupported(), kTfLiteOk); +} + +// Different projections may be blocked differently, since gate/ff1 and linear +// have different input axes. +TEST(MoeDelegateKernelIsSupportedTest, AcceptsMixedBlockCounts) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows * 2, kGateFf1Rows, kLinearRows * 4); + EXPECT_EQ(builder.IsSupported(), kTfLiteOk); +} + +// A scale count that is not a whole number of scales per row cannot describe +// any blocking. +TEST(MoeDelegateKernelIsSupportedTest, RejectsScaleCountNotMultipleOfRows) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows + 1, kGateFf1Rows, kLinearRows); + EXPECT_EQ(builder.IsSupported(), kTfLiteError); +} + +TEST(MoeDelegateKernelIsSupportedTest, RejectsEmptyScaleTensor) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows, kGateFf1Rows, 0); + EXPECT_EQ(builder.IsSupported(), kTfLiteError); +} + +// The exporter now labels the activation "gelu_tanh"; the legacy "gelu" +// spelling has to keep working for already-converted models. +TEST(MoeDelegateKernelIsSupportedTest, AcceptsBothGeluSpellings) { + MoeNodeBuilder tanh_builder; + BuildInt8Moe(tanh_builder, kGateFf1Rows, kGateFf1Rows, kLinearRows, + /*activation=*/"gelu_tanh"); + EXPECT_EQ(tanh_builder.IsSupported(), kTfLiteOk); + + MoeNodeBuilder legacy_builder; + BuildInt8Moe(legacy_builder, kGateFf1Rows, kGateFf1Rows, kLinearRows, + /*activation=*/"gelu"); + EXPECT_EQ(legacy_builder.IsSupported(), kTfLiteOk); +} + +TEST(MoeDelegateKernelIsSupportedTest, RejectsUnknownActivation) { + MoeNodeBuilder builder; + BuildInt8Moe(builder, kGateFf1Rows, kGateFf1Rows, kLinearRows, + /*activation=*/"silu"); + EXPECT_EQ(builder.IsSupported(), kTfLiteError); +} + +} // namespace +} // namespace xnnpack +} // namespace tflite diff --git a/tensorflow/lite/kernels/internal/reference/select.h b/tensorflow/lite/kernels/internal/reference/select.h index 5fe86f174a56ab..408d584c98b20e 100644 --- a/tensorflow/lite/kernels/internal/reference/select.h +++ b/tensorflow/lite/kernels/internal/reference/select.h @@ -109,18 +109,22 @@ void RunSelectOp(const D* cond, const T* x, const T* y, T* output, } else { TFLITE_DCHECK_EQ(cond_stride_0, 1); if (x_stride_0 == 0 && y_stride_0 == 0) { + const T x0 = *x; + const T y0 = *y; for (size_t i = 0; i < output_shape_0; ++i) { - output[i] = cond[i] ? *x : *y; + output[i] = cond[i] ? x0 : y0; } } else if (x_stride_0 == 0) { TFLITE_DCHECK_EQ(y_stride_0, 1); + const T x0 = *x; for (size_t i = 0; i < output_shape_0; ++i) { - output[i] = cond[i] ? *x : y[i]; + output[i] = cond[i] ? x0 : y[i]; } } else if (y_stride_0 == 0) { TFLITE_DCHECK_EQ(x_stride_0, 1); + const T y0 = *y; for (size_t i = 0; i < output_shape_0; ++i) { - output[i] = cond[i] ? x[i] : *y; + output[i] = cond[i] ? x[i] : y0; } } else { TFLITE_DCHECK_EQ(x_stride_0, 1); diff --git a/tensorflow/lite/python/analyzer_wrapper/BUILD b/tensorflow/lite/python/analyzer_wrapper/BUILD index ccf7b77a348340..cd9c91a5f219c2 100644 --- a/tensorflow/lite/python/analyzer_wrapper/BUILD +++ b/tensorflow/lite/python/analyzer_wrapper/BUILD @@ -49,7 +49,6 @@ cc_library( visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/mlir/lite/schema:schema_utils", - "//tensorflow/core:framework", "//tensorflow/core/public:release_version", "//tensorflow/lite:version", "//tensorflow/lite/core:model_builder", diff --git a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc index 8d417c8fde2d6e..abb92666625751 100644 --- a/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc +++ b/tensorflow/lite/toco/graph_transformations/reorder_reshape_transpose.cc @@ -12,7 +12,6 @@ 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 #include #include #include diff --git a/tensorflow/lite/tools/strip_buffers/BUILD b/tensorflow/lite/tools/strip_buffers/BUILD index 723bcb216657d9..1c2ddc1af6b349 100644 --- a/tensorflow/lite/tools/strip_buffers/BUILD +++ b/tensorflow/lite/tools/strip_buffers/BUILD @@ -34,7 +34,6 @@ cc_library( hdrs = ["stripping_lib.h"], copts = tflite_copts(), deps = [ - "//tensorflow/core:tflite_portable_logging", "//tensorflow/lite:framework", "//tensorflow/lite:string", "//tensorflow/lite/core:framework", @@ -54,7 +53,6 @@ cc_binary( linkopts = tflite_linkopts(), deps = [ ":stripping_lib", - "//tensorflow/core:tflite_portable_logging", "//tensorflow/lite:framework", "//tensorflow/lite/core:framework", "//tensorflow/lite/core/c:common", @@ -71,7 +69,6 @@ cc_binary( linkopts = tflite_linkopts(), deps = [ ":stripping_lib", - "//tensorflow/core:tflite_portable_logging", "//tensorflow/lite:framework", "//tensorflow/lite/core:framework", "//tensorflow/lite/core/c:common", diff --git a/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc b/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc index e02cdecef4896a..acd9fcfd257931 100644 --- a/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc +++ b/tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc @@ -20,7 +20,6 @@ limitations under the License. #include "absl/log/log.h" #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/model_builder.h" #include "tensorflow/lite/tools/command_line_flags.h" diff --git a/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc b/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc index 6915075582927a..7f04a2031e46b9 100644 --- a/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc +++ b/tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc @@ -20,7 +20,6 @@ limitations under the License. #include "absl/log/log.h" #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/model_builder.h" #include "tensorflow/lite/tools/command_line_flags.h" diff --git a/tensorflow/lite/tools/strip_buffers/stripping_lib.cc b/tensorflow/lite/tools/strip_buffers/stripping_lib.cc index d29a757e491a21..121705e6343e56 100644 --- a/tensorflow/lite/tools/strip_buffers/stripping_lib.cc +++ b/tensorflow/lite/tools/strip_buffers/stripping_lib.cc @@ -27,7 +27,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "flatbuffers/buffer.h" // from @flatbuffers #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/model_builder.h" #include "tensorflow/lite/schema/schema_generated.h" diff --git a/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py b/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py index 243862993d6b51..bd1b3f3b8b0574 100644 --- a/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py +++ b/tensorflow/python/data/experimental/kernel_tests/service/cross_trainer_cache_test.py @@ -108,8 +108,8 @@ def testConcurrentReaders(self): num_cpus = multiprocessing.cpu_count() cluster = self._create_cluster( num_workers=1, cross_trainer_cache_size_bytes=(num_cpus + 8) * 423) - num_readers = 20 - num_elements = 50 + num_readers = 8 + num_elements = 20 dataset = dataset_ops.Dataset.range(10000000).repeat() datasets = [] diff --git a/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py b/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py index 463b75934f8085..ad85997e0df7e2 100644 --- a/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py +++ b/tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py @@ -4970,6 +4970,78 @@ def testGuardedAssertDoesNotCopyWhenTrue(self): # No copy was performed for the guarded assert self.assertEqual([], guarded_memcpy_nodestat_names) + @test_util.run_in_graph_and_eager_modes + def testAssertRaisesWithRaggedTensorData(self): + rt = ragged_factory_ops.constant([[1], [2]]) + cond = math_ops.reduce_all(math_ops.equal(rt.row_lengths(), 2)) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"\[\[1\], \[2\]\]"): + self.evaluate(control_flow_assert.Assert(cond, [rt])) + + @test_util.run_in_graph_and_eager_modes + def testAssertPassesWithRaggedTensorData(self): + rt = ragged_factory_ops.constant([[1], [2]]) + cond = math_ops.reduce_all(math_ops.equal(rt.row_lengths(), 1)) + self.evaluate(control_flow_assert.Assert(cond, [rt])) + + @test_util.run_in_graph_and_eager_modes + def testAssertRaisesWithDenseAndRaggedTensorData(self): + dense = constant_op.constant([7.5, 8.5]) + rt = ragged_factory_ops.constant([[1, 2, 3, 4], [5]]) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"(?s)7\.5.*\[\[1, 2, 3, 4\], \[5\]\]"): + self.evaluate( + control_flow_assert.Assert(False, [dense, rt], summarize=-1)) + + @test_util.run_in_graph_and_eager_modes + def testAssertSummarizesRaggedTensorData(self): + rt = ragged_factory_ops.constant([[i] for i in range(8)]) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + r"\[\[0\], \[1\], \[2\], \.\.\., \[5\], \[6\], \[7\]\]"): + self.evaluate(control_flow_assert.Assert(False, [rt])) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"\[\[0\], \[1\], \[2\], \[3\], \[4\]"): + self.evaluate(control_flow_assert.Assert(False, [rt], summarize=-2)) + with self.assertRaisesRegex(errors_impl.InvalidArgumentError, + r"\[\[0\], \[1\], \.\.\., \[6\], \[7\]\]"): + self.evaluate( + control_flow_assert.Assert(False, [rt], summarize=np.int64(2))) + with self.assertRaises(errors_impl.InvalidArgumentError): + self.evaluate(control_flow_assert.Assert(False, [rt], summarize=0)) + + @test_util.run_in_graph_and_eager_modes + def testAssertWithUnknownRankRaggedTensorData(self): + + @eager_def_function.function(input_signature=[ + ragged_tensor.RaggedTensorSpec(ragged_rank=1, dtype=dtypes.int32), + tensor_lib.TensorSpec([], dtypes.bool), + ]) + def f(rt, cond): + with ops.control_dependencies([control_flow_assert.Assert(cond, [rt])]): + return array_ops.identity(rt.flat_values) + + rt = ragged_factory_ops.constant([[1, 2], [3]]) + self.assertAllEqual([1, 2, 3], self.evaluate(f(rt, True))) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + r"flat_values=\[1 2 3\], nested_row_splits=\[\[0 2 3\]\]"): + self.evaluate(f(rt, False)) + + @test_util.run_in_graph_and_eager_modes + def testAssertWithRaggedTensorDataCompilesWithXla(self): + if not test_util.is_xla_enabled(): + self.skipTest("Requires XLA.") + + @eager_def_function.function(jit_compile=True) + def f(rt): + cond = math_ops.reduce_all(math_ops.greater(rt.row_lengths(), 0)) + with ops.control_dependencies([control_flow_assert.Assert(cond, [rt])]): + return math_ops.reduce_sum(rt.flat_values) + + self.assertEqual(6, self.evaluate(f(ragged_factory_ops.constant([[1, 2], + [3]])))) + class WhileOpBenchmark(test.Benchmark): """Evaluate the performance of while_loop op.""" diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py index 293d0c67f4b4e4..36760fd2227b80 100644 --- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py @@ -756,6 +756,30 @@ def test_sparse_bincount_input_validation(self): size=size, weights=[])) + def test_sparse_bincount_rank0_fails(self): + # Regression test for the rank-0 guard in SparseBincountOp::Compute + # ("dense_shape must have at least 1 dimension"). Without it the kernel + # reads dense_shape(0) on an empty dense_shape. Shape inference may reject + # the input first, so both messages are accepted. + with self.assertRaisesRegex( + (ValueError, errors.InvalidArgumentError), + "Input must be less than rank 2|" + "dense_shape must have at least 1 dimension", + ): + self.evaluate( + gen_math_ops.sparse_bincount( + indices=constant_op.constant( + [], shape=[1, 0], dtype=dtypes.int64 + ), + values=[0], + dense_shape=constant_op.constant( + [], shape=[0], dtype=dtypes.int64 + ), + size=10, + weights=[], + ) + ) + class RaggedBincountOpTest(test_util.TensorFlowTestCase, parameterized.TestCase): diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 2bcda6a1843e58..675b0563574d5d 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -1123,12 +1123,15 @@ py_library( ":array_ops", ":cond", ":control_flow_ops_gen", + ":control_flow_util", ":logging_ops_gen", ":math_ops_gen", + ":string_ops", "//tensorflow/python/eager:context", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:ops", + "//tensorflow/python/types:internal", "//tensorflow/python/util:dispatch", "//tensorflow/python/util:tf_export", "//tensorflow/python/util:tf_should_use", diff --git a/tensorflow/python/ops/control_flow_assert.py b/tensorflow/python/ops/control_flow_assert.py index 18928fa6f35754..5dd7ee7bff7902 100644 --- a/tensorflow/python/ops/control_flow_assert.py +++ b/tensorflow/python/ops/control_flow_assert.py @@ -20,9 +20,12 @@ from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import cond +from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import gen_control_flow_ops from tensorflow.python.ops import gen_logging_ops from tensorflow.python.ops import gen_math_ops +from tensorflow.python.ops import string_ops +from tensorflow.python.types import internal as internal_types from tensorflow.python.util import dispatch from tensorflow.python.util import tf_should_use from tensorflow.python.util.tf_export import tf_export @@ -57,6 +60,38 @@ def _summarize_eager(tensor, summarize=None): return ", ".join(lst) +def _format_ragged(x, summarize): + """Returns `x`, or a scalar string `Tensor` to pass in its place if ragged. + + Assert only handles dense tensors, so a `RaggedTensor` of known rank is + formatted the way `tf.print` formats it. + + Args: + x: An entry of the `data` argument to `Assert`. + summarize: The `summarize` argument to `Assert`. + """ + if not isinstance(x, internal_types.RaggedTensor): + return x + # Default to Assert's own 3. Ragged formatting accepts only a Python int + # that is -1 (everything) or positive. + summarize = 3 if summarize is None else int(summarize) + if summarize == 0 or control_flow_util.GraphOrParentsInXlaContext( + ops.get_default_graph()): + # Assert shows no entries for summarize=0, and XLA cannot compile the + # formatting ops. + return ops.convert_to_tensor("") + if summarize < 0: + summarize = -1 + if x.shape.rank is None: + # Ragged formatting needs a known rank; fall back to the components. + splits = x.nested_row_splits + template = "tf.RaggedTensor(flat_values={}, nested_row_splits=[%s])" % ( + ", ".join(["{}"] * len(splits))) + return string_ops.string_format( + template, [x.flat_values, *splits], summarize=summarize) + return string_ops.string_format("{}", [x], summarize=summarize) + + # Assert and Print are special symbols in python, so we must # use an upper-case version of them. @tf_export("debugging.Assert", "Assert") @@ -70,8 +105,10 @@ def Assert(condition, data, summarize=None, name=None): Args: condition: The condition to evaluate. - data: The tensors to print out when condition is false. - summarize: Print this many entries of each tensor. + data: The tensors to print out when condition is false. A `RaggedTensor` + of known rank is printed as `tf.print` prints it. + summarize: Print this many entries of each tensor. For a `RaggedTensor`, + the first and last this many entries of each dimension. name: A name for this operation (optional). Returns: @@ -97,6 +134,7 @@ def Assert(condition, data, summarize=None, name=None): """ if context.executing_eagerly(): if not condition: + data = [_format_ragged(x, summarize) for x in data] xs = ops.convert_n_to_tensor(data) data_str = [_summarize_eager(x, summarize) for x in xs] raise errors.InvalidArgumentError( @@ -107,8 +145,13 @@ def Assert(condition, data, summarize=None, name=None): return with ops.name_scope(name, "Assert", [condition, data]) as name: - xs = ops.convert_n_to_tensor(data) - if all(x.dtype in {dtypes.string, dtypes.int32} for x in xs): + # Ragged data cannot go through convert_n_to_tensor, and its formatting + # ops should only run when the assert fails, so it takes the guarded path. + data = list(data) + has_ragged = any(isinstance(x, internal_types.RaggedTensor) for x in data) + if not has_ragged and all( + x.dtype in {dtypes.string, dtypes.int32} + for x in ops.convert_n_to_tensor(data)): # As a simple heuristic, we assume that string and int32 are # on host to avoid the need to use cond. If it is not case, # we will pay the price copying the tensor to host memory. @@ -118,7 +161,8 @@ def Assert(condition, data, summarize=None, name=None): def true_assert(): return gen_logging_ops._assert( # pylint: disable=protected-access - condition, data, summarize, name="Assert") + condition, [_format_ragged(x, summarize) for x in data], + summarize, name="Assert") guarded_assert = cond.cond( condition, diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl index eacb39dbf8a972..bcfe4415614b9f 100644 --- a/tensorflow/workspace2.bzl +++ b/tensorflow/workspace2.bzl @@ -67,6 +67,7 @@ load("@xla//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure") load("@xla//third_party/tensorrt:workspace.bzl", tensorrt = "repo") load("@xla//third_party/triton:workspace.bzl", triton = "repo") load("@xla//third_party/xnnpack:workspace.bzl", xnnpack = "repo") +load("@xla//third_party/xz:workspace.bzl", xz = "repo") load("@xla//tools/def_file_filter:def_file_filter_configure.bzl", "def_file_filter_configure") load("@xla//tools/toolchains:cpus/aarch64/aarch64_compiler_configure.bzl", "aarch64_compiler_configure") load("@xla//tools/toolchains:cpus/arm/arm_compiler_configure.bzl", "arm_compiler_configure") @@ -255,6 +256,7 @@ def _tf_repositories(): ) xnnpack() + xz() # XNNPack dependency. tf_http_archive( diff --git a/third_party/xla/.github/workflows/ci_multi_device.yml b/third_party/xla/.github/workflows/ci_multi_device.yml index 8a8ab8fa70b477..4ea65669e06f76 100644 --- a/third_party/xla/.github/workflows/ci_multi_device.yml +++ b/third_party/xla/.github/workflows/ci_multi_device.yml @@ -69,12 +69,14 @@ jobs: halt-dispatch-input: ${{ inputs.halt-for-connection }} - name: "Run build.py" working-directory: ${{ matrix.job_info.repo }} + env: + PYTHONPATH: ${{ github.workspace }}/openxla/xla run: | # shellcheck disable=SC2193 if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions" else - "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions" + python3 "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions" fi diff --git a/third_party/xla/AGENTS.md b/third_party/xla/AGENTS.md new file mode 100644 index 00000000000000..d3dcfded10316e --- /dev/null +++ b/third_party/xla/AGENTS.md @@ -0,0 +1,11 @@ +# AI Assistant Guidelines for OpenXLA + +Welcome to the OpenXLA (`openxla/xla`) repository. + +* **Core XLA compiler guidelines**: When generating, reviewing, or modifying + code under [`xla/`](xla/), follow the detailed C++, MLIR, testing, and + Bazel guidelines in [`xla/AGENTS.md`](xla/AGENTS.md). +* [`build_tools/`](build_tools/) contains CI and developer tooling, + [`docs/`](docs/) the user-facing documentation, and + [`third_party/`](third_party/) vendored dependencies and Bazel repository + rules. diff --git a/third_party/xla/docs/custom_call.md b/third_party/xla/docs/custom_call.md index 358fdaa26044e5..2356c39dc657cb 100644 --- a/third_party/xla/docs/custom_call.md +++ b/third_party/xla/docs/custom_call.md @@ -309,7 +309,7 @@ The external FFI API in `xla/ffi/api/ffi.h` returns `Error` from `Verify` and ```c++ ABSL_ASSIGN_OR_RETURN(BufferR2 input, - Match("input", buffer, m::Buffer())); + Match("input", buffer, m::Buffer())); ABSL_RETURN_IF_ERROR(Verify( "input", input, m::Buffer().WithDims(expected_rows, m::Dim()))); diff --git a/third_party/xla/opensource_only.files b/third_party/xla/opensource_only.files index 097ae77aa29598..9c227079557dd5 100644 --- a/third_party/xla/opensource_only.files +++ b/third_party/xla/opensource_only.files @@ -160,4 +160,5 @@ xla/third_party/tensorrt/tensorrt/tensorrt_config.py.tpl: xla/third_party/tensorrt/tensorrt_configure.bzl: xla/third_party/tensorrt/workspace.bzl: xla/third_party/xla_googletest_wrapper/include/gmock/gmock.h: +xla/third_party/xz/config/config_generic.h: xla/third_party/zlib.BUILD: diff --git a/third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch b/third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch new file mode 100644 index 00000000000000..821700fa0c8609 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/llvm_cl_989877215.patch @@ -0,0 +1,11 @@ +diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td b/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td +--- a/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td ++++ b/include/cuda_tile/Dialect/CudaTile/IR/Dialect.td +@@ -42,6 +42,7 @@ + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; + ++ let useStrictPropertiesInAssemblyFormat = 0; + let extraClassDeclaration = [{ + template + void addExternalOperations() { diff --git a/third_party/xla/third_party/cuda_tile/patches/series b/third_party/xla/third_party/cuda_tile/patches/series index 9134fcfde4c7bd..46a712b02b5ef8 100644 --- a/third_party/xla/third_party/cuda_tile/patches/series +++ b/third_party/xla/third_party/cuda_tile/patches/series @@ -1,3 +1,4 @@ constructor.patch unused_variable.patch pow_op_alias.patch +llvm_cl_989877215.patch diff --git a/third_party/xla/third_party/extensions/llvm.bzl b/third_party/xla/third_party/extensions/llvm.bzl index 8d3262d8cb498f..b3ca0e0c658fdc 100644 --- a/third_party/xla/third_party/extensions/llvm.bzl +++ b/third_party/xla/third_party/extensions/llvm.bzl @@ -19,6 +19,7 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure") load("@llvm-raw//utils/bazel:linux_uapi.bzl", "linux_uapi_setup") load("//third_party:repo.bzl", "tf_mirror_urls") +load("//third_party/xz:workspace.bzl", xz = "repo") _PYYAML_CONTENT = """\ load("@rules_python//python:defs.bzl", "py_library") @@ -64,6 +65,7 @@ _llvm_zstd_compat = repository_rule( def _llvm_extension_impl(mctx): # @unused _llvm_zlib_compat(name = "llvm_zlib") _llvm_zstd_compat(name = "llvm_zstd") + xz() linux_uapi_setup(name = "linux_uapi") http_archive( name = "pyyaml", diff --git a/third_party/xla/third_party/llvm/toolchains.patch b/third_party/xla/third_party/llvm/toolchains.patch index a8c15b85d9235a..1170c7610ab573 100644 --- a/third_party/xla/third_party/llvm/toolchains.patch +++ b/third_party/xla/third_party/llvm/toolchains.patch @@ -13,13 +13,13 @@ # limitations under the License. # ============================================================================== diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -index b683f86ca386..70819c622402 100644 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -59,6 +59,36 @@ exports_files([ +@@ -57,6 +57,36 @@ + # This one is needed for building and vendoring out lldb from off tree. "utils/lldbDataFormatters.py", ]) - ++ +config_setting( + name = "macos_arm64", + constraint_values = [ @@ -49,20 +49,9 @@ index b683f86ca386..70819c622402 100644 + "@platforms//cpu:x86_64", + ], +) -+ + config_setting( name = "darwin_arm64", - constraint_values = [ -@@ -533,7 +563,7 @@ cc_library( - # We unconditionally depend on the custom LLVM lzma wrapper. This only - # provides the configuration define unless lzma is enabled, in which - # case it also links the system liblzma. -- "//third-party:lzma", -+ # "//third-party:lzma", - "//libc:shared_math_headers_for_apfloat", - ], - ) - diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl --- a/utils/bazel/llvm-project-overlay/llvm/config.bzl +++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl diff --git a/third_party/xla/third_party/triton/common/series.bzl b/third_party/xla/third_party/triton/common/series.bzl index 3c729811758214..e5a4592a9bec1b 100644 --- a/third_party/xla/third_party/triton/common/series.bzl +++ b/third_party/xla/third_party/triton/common/series.bzl @@ -64,5 +64,6 @@ common_patch_list = [ "//third_party/triton:common/llvm_cl979526570.patch", "//third_party/triton:common/llvm_cl983115707.patch", "//third_party/triton:common/llvm_cl986804242.patch", + "//third_party/triton:common/useStrictPropertiesInAssemblyFormat.patch", # Add new patches just above this line ] diff --git a/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch b/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch new file mode 100644 index 00000000000000..83639d1da45ea9 --- /dev/null +++ b/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch @@ -0,0 +1,114 @@ +Prevents issues due to https://github.com/llvm/llvm-project/commit/c495a9b796c46456e523a5eeb3690415ba99f5db. +This is a temporary fix to unblock the LLVM integration, the opt-out will eventually be removed. +See b/567092262 for more details. + +--- a/include/triton/Dialect/Gluon/IR/GluonDialect.td ++++ b/include/triton/Dialect/Gluon/IR/GluonDialect.td +@@ -16,6 +16,7 @@ + "mlir::gpu::GPUDialect", + ]; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/include/triton/Dialect/Triton/IR/TritonDialect.td ++++ b/include/triton/Dialect/Triton/IR/TritonDialect.td +@@ -51,6 +51,7 @@ + + let hasConstantMaterializer = 1; + let useDefaultTypePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + include "triton/Dialect/Triton/IR/TritonTypes.td" + +--- a/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td ++++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td +@@ -35,6 +35,7 @@ + + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td ++++ b/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td +@@ -6,6 +6,7 @@ + def TritonInstrument_Dialect : Dialect { + let name = "tti"; + let cppNamespace = "::mlir::triton::instrument"; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif // TRITONINSTRUMENT_DIALECT + +--- a/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td ++++ b/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td +@@ -43,6 +43,7 @@ + + let useDefaultAttributePrinterParser = 1; + let useDefaultTypePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td ++++ b/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td +@@ -37,6 +37,7 @@ + let dependentDialects = ["triton::TritonDialect"]; + + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td ++++ b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td +@@ -35,6 +35,7 @@ + let dependentDialects = [ + "mlir::LLVM::LLVMDialect" + ]; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/nvidia/include/Dialect/NVWS/IR/NVWSDialect.td ++++ b/third_party/nvidia/include/Dialect/NVWS/IR/NVWSDialect.td +@@ -39,6 +39,7 @@ + + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/proton/Dialect/include/Dialect/Proton/IR/ProtonDialect.td ++++ b/third_party/proton/Dialect/include/Dialect/Proton/IR/ProtonDialect.td +@@ -13,6 +13,7 @@ + }]; + + let dependentDialects = []; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif + +--- a/third_party/proton/Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUDialect.td ++++ b/third_party/proton/Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUDialect.td +@@ -22,6 +22,7 @@ + + let useDefaultTypePrinterParser = 1; + let useDefaultAttributePrinterParser = 1; ++ let useStrictPropertiesInAssemblyFormat = 0; + } + + #endif // PROTONGPU_DIALECT + diff --git a/third_party/xla/third_party/xz/BUILD.bazel b/third_party/xla/third_party/xz/BUILD.bazel new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/third_party/xla/third_party/xz/config/BUILD.bazel b/third_party/xla/third_party/xz/config/BUILD.bazel new file mode 100644 index 00000000000000..d0efcb1c3bd83c --- /dev/null +++ b/third_party/xla/third_party/xz/config/BUILD.bazel @@ -0,0 +1 @@ +exports_files(["config_generic.h"]) diff --git a/third_party/xla/third_party/xz/config/config_generic.h b/third_party/xla/third_party/xz/config/config_generic.h new file mode 100644 index 00000000000000..78807e812e8722 --- /dev/null +++ b/third_party/xla/third_party/xz/config/config_generic.h @@ -0,0 +1,557 @@ +/* Copyright 2020 The TensorStore 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. +==============================================================================*/ + +/* config.h. Generated from config.h.in by configure. */ +/* config.h.in. Generated from configure.ac by autoheader. */ + +/* Define if building universal (internal helper macro) */ +/* #undef AC_APPLE_UNIVERSAL_BUILD */ + +/* How many MiB of RAM to assume if the real amount cannot be determined. */ +#define ASSUME_RAM 128 + +/* Define to 1 if translation of program messages to the user's native + language is requested. */ +/* #undef ENABLE_NLS */ + +#ifndef _WIN32 +/* Define to 1 if bswap_16 is available. */ +#define HAVE_BSWAP_16 1 + +/* Define to 1 if bswap_32 is available. */ +#define HAVE_BSWAP_32 1 + +/* Define to 1 if bswap_64 is available. */ +#define HAVE_BSWAP_64 1 +#endif + +/* Define to 1 if you have the header file. */ +#if defined(__APPLE__) || defined(_WIN32) +// byteswap.h does not exist in Apple SDKs. +// This #define must be entirely eliminated (not 0), as usages are checking +// for it being defined. +#undef HAVE_BYTESWAP_H +#else +#define HAVE_BYTESWAP_H 1 +#endif + +/* Define to 1 if Capsicum is available. */ +#undef HAVE_CAPSICUM + +/* Define to 1 if the system has the type `CC_SHA256_CTX'. */ +#if defined(__APPLE__) +#define HAVE_CC_SHA256_CTX 1 +#endif + +/* Define to 1 if you have the `CC_SHA256_Init' function. */ +#if defined(__APPLE__) +#define HAVE_CC_SHA256_INIT 1 +#endif + +/* Define to 1 if you have the MacOS X function CFLocaleCopyCurrent in the + CoreFoundation framework. */ +/* #undef HAVE_CFLOCALECOPYCURRENT */ + +/* Define to 1 if you have the MacOS X function CFPreferencesCopyAppValue in + the CoreFoundation framework. */ +/* #undef HAVE_CFPREFERENCESCOPYAPPVALUE */ + +/* Define to 1 if crc32 integrity check is enabled. */ +#define HAVE_CHECK_CRC32 1 + +/* Define to 1 if crc64 integrity check is enabled. */ +#define HAVE_CHECK_CRC64 1 + +/* Define to 1 if sha256 integrity check is enabled. */ +#define HAVE_CHECK_SHA256 1 + +/* Define to 1 if you have the `clock_gettime' function. */ +#if !defined(__APPLE__) +#define HAVE_CLOCK_GETTIME 1 +#endif + +/* Define to 1 if you have the header file. */ +#if defined(__APPLE__) +#define HAVE_COMMONCRYPTO_COMMONDIGEST_H 1 +#endif + +/* Define if the GNU dcgettext() function is already present or preinstalled. + */ +/* #undef HAVE_DCGETTEXT */ + +/* Define to 1 if you have the declaration of `CLOCK_MONOTONIC', and to 0 if + you don't. */ +#if !defined(__APPLE__) +#define HAVE_DECL_CLOCK_MONOTONIC 1 +#endif + +/* Define to 1 if you have the declaration of `program_invocation_name', and + to 0 if you don't. */ +#if defined(__ANDROID__) || defined(__APPLE__) || defined(_WIN32) +#define HAVE_DECL_PROGRAM_INVOCATION_NAME 0 +#else +#define HAVE_DECL_PROGRAM_INVOCATION_NAME 1 +#endif + +/* Define to 1 if any of HAVE_DECODER_foo have been defined. */ +#define HAVE_DECODERS 1 + +/* Define to 1 if arm decoder is enabled. */ +#define HAVE_DECODER_ARM 1 + +/* Define to 1 if armthumb decoder is enabled. */ +#define HAVE_DECODER_ARMTHUMB 1 + +/* Define to 1 if delta decoder is enabled. */ +#define HAVE_DECODER_DELTA 1 + +/* Define to 1 if ia64 decoder is enabled. */ +#define HAVE_DECODER_IA64 1 + +/* Define to 1 if lzma1 decoder is enabled. */ +#define HAVE_DECODER_LZMA1 1 + +/* Define to 1 if lzma2 decoder is enabled. */ +#define HAVE_DECODER_LZMA2 1 + +/* Define to 1 if powerpc decoder is enabled. */ +#define HAVE_DECODER_POWERPC 1 + +/* Define to 1 if sparc decoder is enabled. */ +#define HAVE_DECODER_SPARC 1 + +/* Define to 1 if x86 decoder is enabled. */ +#define HAVE_DECODER_X86 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_DLFCN_H 1 + +/* Define to 1 if any of HAVE_ENCODER_foo have been defined. */ +#define HAVE_ENCODERS 1 + +/* Define to 1 if arm encoder is enabled. */ +#define HAVE_ENCODER_ARM 1 + +/* Define to 1 if armthumb encoder is enabled. */ +#define HAVE_ENCODER_ARMTHUMB 1 + +/* Define to 1 if delta encoder is enabled. */ +#define HAVE_ENCODER_DELTA 1 + +/* Define to 1 if ia64 encoder is enabled. */ +#define HAVE_ENCODER_IA64 1 + +/* Define to 1 if lzma1 encoder is enabled. */ +#define HAVE_ENCODER_LZMA1 1 + +/* Define to 1 if lzma2 encoder is enabled. */ +#define HAVE_ENCODER_LZMA2 1 + +/* Define to 1 if powerpc encoder is enabled. */ +#define HAVE_ENCODER_POWERPC 1 + +/* Define to 1 if sparc encoder is enabled. */ +#define HAVE_ENCODER_SPARC 1 + +/* Define to 1 if x86 encoder is enabled. */ +#define HAVE_ENCODER_X86 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_FCNTL_H 1 + +/* Define to 1 if you have the `futimens' function. */ +#define HAVE_FUTIMENS 1 + +/* Define to 1 if you have the `futimes' function. */ +/* #undef HAVE_FUTIMES */ + +/* Define to 1 if you have the `futimesat' function. */ +/* #undef HAVE_FUTIMESAT */ + +/* Define to 1 if you have the header file. */ +#define HAVE_GETOPT_H 1 + +/* Define to 1 if you have the `getopt_long' function. */ +#define HAVE_GETOPT_LONG 1 + +/* Define if the GNU gettext() function is already present or preinstalled. */ +/* #undef HAVE_GETTEXT */ + +/* Define if you have the iconv() function and it works. */ +/* #undef HAVE_ICONV */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_IMMINTRIN_H */ + +/* Define to 1 if you have the header file. */ +#define HAVE_INTTYPES_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_LIMITS_H 1 + +/* Define to 1 if mbrtowc and mbstate_t are properly declared. */ +#define HAVE_MBRTOWC 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_MEMORY_H 1 + +/* Define to 1 to enable bt2 match finder. */ +#define HAVE_MF_BT2 1 + +/* Define to 1 to enable bt3 match finder. */ +#define HAVE_MF_BT3 1 + +/* Define to 1 to enable bt4 match finder. */ +#define HAVE_MF_BT4 1 + +/* Define to 1 to enable hc3 match finder. */ +#define HAVE_MF_HC3 1 + +/* Define to 1 to enable hc4 match finder. */ +#define HAVE_MF_HC4 1 + +/* Define to 1 if getopt.h declares extern int optreset. */ +/* #undef HAVE_OPTRESET */ + +/* Define to 1 if you have the `posix_fadvise' function. */ +#if !defined(__APPLE__) +#define HAVE_POSIX_FADVISE 1 +#endif + +/* Define to 1 if you have the `pthread_condattr_setclock' function. */ +#if !defined(__APPLE__) && !defined(__ANDROID__) +#define HAVE_PTHREAD_CONDATTR_SETCLOCK 1 +#endif + +/* Have PTHREAD_PRIO_INHERIT. */ +#define HAVE_PTHREAD_PRIO_INHERIT 1 + +/* Define to 1 if you have the `SHA256Init' function. */ +/* #undef HAVE_SHA256INIT */ + +/* Define to 1 if the system has the type `SHA256_CTX'. */ +/* #undef HAVE_SHA256_CTX */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SHA256_H */ + +/* Define to 1 if you have the `SHA256_Init' function. */ +/* #undef HAVE_SHA256_INIT */ + +/* Define to 1 if the system has the type `SHA2_CTX'. */ +/* #undef HAVE_SHA2_CTX */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SHA2_H */ + +/* Define to 1 if optimizing for size. */ +#undef HAVE_SMALL + +/* Define to 1 if stdbool.h conforms to C99. */ +#define HAVE_STDBOOL_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_STDINT_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_STDLIB_H 1 + +/* Define to 1 if you have the header file. */ +/* #define HAVE_STRINGS_H 1 */ + +/* Define to 1 if you have the header file. */ +#define HAVE_STRING_H 1 + +/* Define to 1 if `st_atimensec' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_ATIMENSEC */ + +/* Define to 1 if `st_atimespec.tv_nsec' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_ATIMESPEC_TV_NSEC */ + +/* Define to 1 if `st_atim.st__tim.tv_nsec' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_ATIM_ST__TIM_TV_NSEC */ + +/* Define to 1 if `st_atim.tv_nsec' is a member of `struct stat'. */ +#define HAVE_STRUCT_STAT_ST_ATIM_TV_NSEC 1 + +/* Define to 1 if `st_uatime' is a member of `struct stat'. */ +/* #undef HAVE_STRUCT_STAT_ST_UATIME */ + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SYS_BYTEORDER_H */ + +/* Define to 1 if you have the header file. */ +#undef HAVE_SYS_CAPSICUM_H + +/* Define to 1 if you have the header file. */ +/* #undef HAVE_SYS_ENDIAN_H */ + +#ifndef _WIN32 +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_PARAM_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_STAT_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_TIME_H 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_SYS_TYPES_H 1 +#endif + +/* Define to 1 if the system has the type `uintptr_t'. */ +#define HAVE_UINTPTR_T 1 + +/* Define to 1 if you have the header file. */ +#define HAVE_UNISTD_H 1 + +/* Define to 1 if you have the `utime' function. */ +/* #undef HAVE_UTIME */ + +/* Define to 1 if you have the `utimes' function. */ +/* #undef HAVE_UTIMES */ + +/* Define to 1 or 0, depending whether the compiler supports simple visibility + declarations. */ +#define HAVE_VISIBILITY 0 + +#ifndef _WIN32 +/* Define to 1 if you have the `wcwidth' function. */ +#define HAVE_WCWIDTH 1 +#endif + +/* Define to 1 if the system has the type `_Bool'. */ +#define HAVE__BOOL 1 + +/* Define to 1 if you have the `_futime' function. */ +#undef HAVE__FUTIME + +/* Define to 1 if _mm_movemask_epi8 is available. */ +/* #undef HAVE__MM_MOVEMASK_EPI8 */ + +/* Define to the sub-directory where libtool stores uninstalled libraries. */ +#define LT_OBJDIR ".libs/" + +#ifdef _WIN32 + +/* Define to 1 when using Windows Vista compatible threads. This uses features + that are not available on Windows XP. */ +#define MYTHREAD_VISTA 1 + +#else + +/* Define to 1 when using POSIX threads (pthreads). */ +#define MYTHREAD_POSIX 1 + +#endif + +/* Define to 1 when using Windows 95 (and thus XP) compatible threads. This + avoids use of features that were added in Windows Vista. */ +/* #undef MYTHREAD_WIN95 */ + +/* Define to 1 to disable debugging code. */ +#define NDEBUG 1 + +/* Name of package */ +#define PACKAGE "xz" + +/* Define to the address where bug reports for this package should be sent. */ +#define PACKAGE_BUGREPORT "lasse.collin@tukaani.org" + +/* Define to the full name of this package. */ +#define PACKAGE_NAME "XZ Utils" + +/* Define to the full name and version of this package. */ +#define PACKAGE_STRING "XZ Utils 5.2.4" + +/* Define to the one symbol short name of this package. */ +#define PACKAGE_TARNAME "xz" + +/* Define to the home page for this package. */ +#define PACKAGE_URL "http://tukaani.org/xz/" + +/* Define to the version of this package. */ +#define PACKAGE_VERSION "5.2.4" + +/* Define to necessary symbol if this constant uses a non-standard name on + your system. */ +/* #undef PTHREAD_CREATE_JOINABLE */ + +/* The size of `size_t', as computed by sizeof. */ +#define SIZEOF_SIZE_T 8 + +/* Define to 1 if you have the ANSI C header files. */ +#define STDC_HEADERS 1 + +/* Define to 1 if the number of available CPU cores can be detected with + cpuset(2). */ +/* #undef TUKLIB_CPUCORES_CPUSET */ + +/* Define to 1 if the number of available CPU cores can be detected with + pstat_getdynamic(). */ +/* #undef TUKLIB_CPUCORES_PSTAT_GETDYNAMIC */ + +/* Define to 1 if the number of available CPU cores can be detected with + sysconf(_SC_NPROCESSORS_ONLN) or sysconf(_SC_NPROC_ONLN). */ +#define TUKLIB_CPUCORES_SYSCONF 1 + +/* Define to 1 if the number of available CPU cores can be detected with + sysctl(). */ +/* #undef TUKLIB_CPUCORES_SYSCTL */ + +/* Define to 1 if the system supports fast unaligned access to 16-bit and + 32-bit integers. */ +#define TUKLIB_FAST_UNALIGNED_ACCESS 1 + +/* Define to 1 if the amount of physical memory can be detected with + _system_configuration.physmem. */ +/* #undef TUKLIB_PHYSMEM_AIX */ + +/* Define to 1 if the amount of physical memory can be detected with + getinvent_r(). */ +/* #undef TUKLIB_PHYSMEM_GETINVENT_R */ + +/* Define to 1 if the amount of physical memory can be detected with + getsysinfo(). */ +/* #undef TUKLIB_PHYSMEM_GETSYSINFO */ + +/* Define to 1 if the amount of physical memory can be detected with + pstat_getstatic(). */ +/* #undef TUKLIB_PHYSMEM_PSTAT_GETSTATIC */ + +/* Define to 1 if the amount of physical memory can be detected with + sysconf(_SC_PAGESIZE) and sysconf(_SC_PHYS_PAGES). */ +#if defined(__APPLE__) +// Apple SDKs (either iOS or OSX) do not define _SC_PHYS_PAGES. +// This #define must be entirely eliminated (not 0), as usages are checking +// for it being defined. +#undef TUKLIB_PHYSMEM_SYSCONF +#else +#define TUKLIB_PHYSMEM_SYSCONF 1 +#endif + +/* Define to 1 if the amount of physical memory can be detected with sysctl(). + */ +/* #undef TUKLIB_PHYSMEM_SYSCTL */ + +/* Define to 1 if the amount of physical memory can be detected with Linux + sysinfo(). */ +/* #undef TUKLIB_PHYSMEM_SYSINFO */ + +/* Enable extensions on AIX 3, Interix. */ +#ifndef _ALL_SOURCE +# define _ALL_SOURCE 1 +#endif +/* Enable GNU extensions on systems that have them. */ +#ifndef _GNU_SOURCE +# define _GNU_SOURCE 1 +#endif +/* Enable threading extensions on Solaris. */ +#ifndef _POSIX_PTHREAD_SEMANTICS +# define _POSIX_PTHREAD_SEMANTICS 1 +#endif +/* Enable extensions on HP NonStop. */ +#ifndef _TANDEM_SOURCE +# define _TANDEM_SOURCE 1 +#endif +/* Enable general extensions on Solaris. */ +#ifndef __EXTENSIONS__ +# define __EXTENSIONS__ 1 +#endif + + +/* Version number of package */ +#define VERSION "5.2.4" + +/* Define WORDS_BIGENDIAN to 1 if your processor stores words with the most + significant byte first (like Motorola and SPARC, unlike Intel). */ +#if defined AC_APPLE_UNIVERSAL_BUILD +# if defined __BIG_ENDIAN__ +# define WORDS_BIGENDIAN 1 +# endif +#else +# ifndef WORDS_BIGENDIAN +/* # undef WORDS_BIGENDIAN */ +# endif +#endif + +/* Enable large inode numbers on Mac OS X 10.5. */ +#ifndef _DARWIN_USE_64_BIT_INODE +# define _DARWIN_USE_64_BIT_INODE 1 +#endif + +/* Number of bits in a file offset, on hosts where this is settable. */ +/* #undef _FILE_OFFSET_BITS */ + +/* Define for large files, on AIX-style hosts. */ +/* #undef _LARGE_FILES */ + +/* Define to 1 if on MINIX. */ +/* #undef _MINIX */ + +/* Define to 2 if the system does not provide POSIX.1 features except with + this defined. */ +/* #undef _POSIX_1_SOURCE */ + +/* Define to 1 if you need to in order for `stat' and other things to work. */ +/* #undef _POSIX_SOURCE */ + +/* Define for Solaris 2.5.1 so the uint32_t typedef from , + , or is not used. If the typedef were allowed, the + #define below would cause a syntax error. */ +/* #undef _UINT32_T */ + +/* Define for Solaris 2.5.1 so the uint64_t typedef from , + , or is not used. If the typedef were allowed, the + #define below would cause a syntax error. */ +/* #undef _UINT64_T */ + +/* Define for Solaris 2.5.1 so the uint8_t typedef from , + , or is not used. If the typedef were allowed, the + #define below would cause a syntax error. */ +/* #undef _UINT8_T */ + +/* Define to rpl_ if the getopt replacement functions and variables should be + used. */ +/* #undef __GETOPT_PREFIX */ + +/* Define to the type of a signed integer type of width exactly 32 bits if + such a type exists and the standard includes do not define it. */ +/* #undef int32_t */ + +/* Define to the type of a signed integer type of width exactly 64 bits if + such a type exists and the standard includes do not define it. */ +/* #undef int64_t */ + +/* Define to the type of an unsigned integer type of width exactly 16 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint16_t */ + +/* Define to the type of an unsigned integer type of width exactly 32 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint32_t */ + +/* Define to the type of an unsigned integer type of width exactly 64 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint64_t */ + +/* Define to the type of an unsigned integer type of width exactly 8 bits if + such a type exists and the standard includes do not define it. */ +/* #undef uint8_t */ + +/* Define to the type of an unsigned integer type wide enough to hold a + pointer, if such a type exists, and if the system does not define it. */ +/* #undef uintptr_t */ diff --git a/third_party/xla/third_party/xz/patches/BUILD.bazel b/third_party/xla/third_party/xz/patches/BUILD.bazel new file mode 100644 index 00000000000000..1d4b626f2c2ffb --- /dev/null +++ b/third_party/xla/third_party/xz/patches/BUILD.bazel @@ -0,0 +1 @@ +exports_files(["remove_have_config.diff"]) diff --git a/third_party/xla/third_party/xz/patches/remove_have_config.diff b/third_party/xla/third_party/xz/patches/remove_have_config.diff new file mode 100644 index 00000000000000..8d6612736c8df6 --- /dev/null +++ b/third_party/xla/third_party/xz/patches/remove_have_config.diff @@ -0,0 +1,14 @@ +diff -Nur a/src/common/sysdefs.h b/src/common/sysdefs.h +--- a/src/common/sysdefs.h 2025-04-03 11:34:58.000000000 +0000 ++++ b/src/common/sysdefs.h 2025-06-20 23:44:13.985352903 +0000 +@@ -19,9 +19,7 @@ + // Includes // + ////////////// + +-#ifdef HAVE_CONFIG_H +-# include +-#endif ++#include "lzma_config/config.h" + + // Choose if MinGW-w64's stdio replacement functions should be used. + // The default has varied slightly in the past so it's clearest to always diff --git a/third_party/xla/third_party/xz/workspace.bzl b/third_party/xla/third_party/xz/workspace.bzl new file mode 100644 index 00000000000000..22e73727d36198 --- /dev/null +++ b/third_party/xla/third_party/xz/workspace.bzl @@ -0,0 +1,37 @@ +# Copyright 2020 The TensorStore 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. + +# buildifier: disable=module-docstring + +load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe") +load( + "//third_party:repo.bzl", + "tf_http_archive", + "tf_mirror_urls", +) + +def repo(): + maybe( + tf_http_archive, + # The repo name is the same as the one used in llvm_project to avoid the conflict with `xz` + # repository in rules_ml_toolchain. + name = "llvm_xz", + strip_prefix = "xz-5.8.3", + urls = tf_mirror_urls("https://github.com/tukaani-project/xz/releases/download/v5.8.3/xz-5.8.3.tar.gz"), + sha256 = "3d3a1b973af218114f4f889bbaa2f4c037deaae0c8e815eec381c3d546b974a0", + patch_file = [ + "//third_party/xz/patches:remove_have_config.diff", + ], + build_file = "//third_party/xz:xz.BUILD.bazel", + ) diff --git a/third_party/xla/third_party/xz/xz.BUILD.bazel b/third_party/xla/third_party/xz/xz.BUILD.bazel new file mode 100644 index 00000000000000..4aacd302a267c7 --- /dev/null +++ b/third_party/xla/third_party/xz/xz.BUILD.bazel @@ -0,0 +1,172 @@ +load("@bazel_skylib//rules:copy_file.bzl", "copy_file") +load("@rules_cc//cc:cc_library.bzl", "cc_library") + +package(default_visibility = ["//visibility:private"]) + +licenses(["unencumbered"]) + +exports_files(["COPYING"]) + +cc_library( + name = "lzma", + srcs = glob( + include = [ + "src/**/*.h", + "src/liblzma/**/*.c", + ], + exclude = [ + "src/**/*_tablegen.c", + "src/liblzma/check/crc_clmul_consts_gen.c", + "src/liblzma/check/sha256.c", + "src/liblzma/check/crc*_small.c", + "src/liblzma/common/stream_decoder_mt.c", + "src/liblzma/common/stream_encoder_mt.c", + ], + ) + select({ + "@platforms//os:osx": [], + "@platforms//os:ios": [], + "//conditions:default": [ + "src/liblzma/check/sha256.c", + ], + }) + select({ + "@platforms//os:none": [], + "//conditions:default": [ + "src/common/tuklib_cpucores.c", + "src/common/tuklib_physmem.c", + "src/liblzma/common/stream_decoder_mt.c", + "src/liblzma/common/stream_encoder_mt.c", + ], + }), + hdrs = glob(["src/liblzma/api/**/*.h"]), + copts = select({ + "@platforms//os:windows": [], + "//conditions:default": ["-std=c99"], + }), + defines = select({ + "@platforms//os:windows": ["LZMA_API_STATIC"], + "//conditions:default": [], + }), + implementation_deps = [ + ":config", + ":src_common", + ":src_liblzma", + ":src_liblzma_check", + ":src_liblzma_common", + ":src_liblzma_delta", + ":src_liblzma_lz", + ":src_liblzma_lzma", + ":src_liblzma_rangecoder", + ":src_liblzma_simple", + ], + linkopts = select({ + "@platforms//os:android": [], + "@platforms//os:windows": [], + "@platforms//os:none": [], + "//conditions:default": ["-lpthread"], + }), + linkstatic = select({ + "@platforms//os:windows": True, + "//conditions:default": False, + }), + strip_include_prefix = "src/liblzma/api", + visibility = ["//visibility:public"], +) + +# Generate config.h from config_generic.h +copy_file( + name = "config_copy", + src = "@xla//third_party/xz/config:config_generic.h", + out = "lzma_config/config.h", +) + +cc_library( + name = "config", + hdrs = ["lzma_config/config.h"], +) + +cc_library( + name = "src_common", + srcs = select({ + "@platforms//os:none": [], + "//conditions:default": [ + "src/common/tuklib_exit.c", + "src/common/tuklib_progname.c", + ], + }), + hdrs = glob(["src/common/*.h"]), + defines = select({ + "@platforms//os:windows": ["LZMA_API_STATIC"], + "//conditions:default": [], + }), + local_defines = select({ + "@platforms//os:windows": ["TUKLIB_GETTEXT=0"], + "//conditions:default": [], + }), + strip_include_prefix = "src/common", + deps = [ + ":config", + ":src_liblzma", + ":src_liblzma_api", + ":src_liblzma_common", + ], +) + +cc_library( + name = "src_liblzma", + hdrs = glob([ + "src/liblzma/common/*.h", + "src/liblzma/lzma/*.h", + ]), + strip_include_prefix = "src/liblzma", +) + +cc_library( + name = "src_liblzma_api", + hdrs = glob( + include = ["src/liblzma/api/**/*.h"], + exclude = ["src/liblzma/api/lzma.h"], + ), + strip_include_prefix = "src/liblzma/api", +) + +cc_library( + name = "src_liblzma_check", + hdrs = glob(["src/liblzma/check/*.h"]), + strip_include_prefix = "src/liblzma/check", +) + +cc_library( + name = "src_liblzma_common", + hdrs = glob(["src/liblzma/common/*.h"]), + strip_include_prefix = "src/liblzma/common", +) + +cc_library( + name = "src_liblzma_delta", + hdrs = glob(["src/liblzma/delta/*.h"]), + strip_include_prefix = "src/liblzma/delta", +) + +cc_library( + name = "src_liblzma_lz", + hdrs = glob(["src/liblzma/lz/*.h"]), + strip_include_prefix = "src/liblzma/lz", +) + +cc_library( + name = "src_liblzma_lzma", + hdrs = glob(["src/liblzma/lzma/*.h"]), + strip_include_prefix = "src/liblzma/lzma", +) + +cc_library( + name = "src_liblzma_rangecoder", + hdrs = glob(["src/liblzma/rangecoder/*.h"]), + strip_include_prefix = "src/liblzma/rangecoder", +) + +cc_library( + name = "src_liblzma_simple", + hdrs = glob(["src/liblzma/simple/*.h"]), + strip_include_prefix = "src/liblzma/simple", +) diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl index 8a99921e4d611f..00a196dcefbfa2 100644 --- a/third_party/xla/workspace2.bzl +++ b/third_party/xla/workspace2.bzl @@ -90,6 +90,7 @@ load("//third_party/transformer_engine:workspace.bzl", transformer_engine = "rep load("//third_party/triton:workspace.bzl", triton = "repo") load("//third_party/uv:workspace.bzl", uv = "repo") load("//third_party/xnnpack:workspace.bzl", xnnpack = "repo") +load("//third_party/xz:workspace.bzl", xz = "repo") load("//tools/def_file_filter:def_file_filter_configure.bzl", "def_file_filter_configure") load("//tools/toolchains:cpus/aarch64/aarch64_compiler_configure.bzl", "aarch64_compiler_configure") load("//tools/toolchains:cpus/arm/arm_compiler_configure.bzl", "arm_compiler_configure") @@ -157,6 +158,7 @@ def _initialize_third_party(): triton() uv() xnnpack() + xz() cuda_tile() cutlass() cutlass_cutedsl_runtime() diff --git a/third_party/xla/xla/AGENTS.md b/third_party/xla/xla/AGENTS.md index 4b59d200fc7bdf..1d945983879635 100644 --- a/third_party/xla/xla/AGENTS.md +++ b/third_party/xla/xla/AGENTS.md @@ -1,8 +1,8 @@ # AI Assistant Guidelines for OpenXLA Development This document provides guidelines for AI code assistants when generating, -suggesting, or modifying code within the -`third_party/tensorflow/compiler/xla` (OpenXLA) codebase. +suggesting, or modifying code within the OpenXLA codebase (everything under +`xla/` in this repository). ## General Context @@ -10,12 +10,27 @@ suggesting, or modifying code within the * **Code Quality:** Adhere to [Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html) and OpenXLA-specific conventions. * **Portability:** This code is open-sourced and runs on a number of host platforms (e.g. Linux, Windows, etc.) +## Repository Layout and Build + +* Source lives under `xla/`. Include paths are relative to the repository + root, e.g. `#include "xla/hlo/ir/hlo_module.h"`. +* Bazel labels use the same prefix, e.g. `//xla/hlo/ir:hlo_module`. +* Build and test with Bazel, e.g. `bazel test //xla/service:hlo_verifier_test`. +* Common external dependencies and their labels: + `@com_google_absl//absl/...`, `@llvm-project//mlir/...`, + `@llvm-project//llvm/...`, `@com_google_googletest//:gtest_main`, + `@com_google_protobuf//:protobuf`. TSL (the platform layer) lives in-tree + under `//xla/tsl/...`. +* Test macros: `xla_cc_test` from `//xla:xla.default.bzl` (hardware- + independent unit tests) and `xla_test` from `//xla/tests:build_defs.bzl` + (tests that run on configured hardware backends). + ## Coding Guidelines for AI Assistance 1. **Error Handling (`absl::Status`, `absl::StatusOr`)**: * **Always** use `absl::Status` or `absl::StatusOr` for functions that can encounter recoverable errors. * **Macros**: - * Use header `tsl/platform/status_macros.h`. + * Use header `absl/status/status_macros.h`. * Use `ABSL_RETURN_IF_ERROR` for error propagation. * Use `ABSL_ASSIGN_OR_RETURN` for `StatusOr` assignments. * **Safely access `StatusOr` values**: Check `.ok()` before accessing. @@ -35,13 +50,14 @@ suggesting, or modifying code within the return absl::OkStatus(); } ``` + 3. **Decision making**: * **Avoid `bool`** for returning a decision to do or not to do something. * **Prefer `Decision`**: - * Located in `third_party/tensorflow/compiler/xla/service/decision.h` + * Located in `xla/service/decision.h` * Example: ```cpp - #include "third_party/tensorflow/compiler/xla/service/decision.h" + #include "xla/service/decision.h" using AutotunerDecision = Decision; @@ -50,7 +66,7 @@ suggesting, or modifying code within the return AutotunerDecision::Forbid("Cublas autotuning was explicitly disabled"); } - voud AutotuneCublas(const AutotunerDecision& decision) { + void AutotuneCublas(const AutotunerDecision& decision) { if (decision.IsForbidden()) { return; } @@ -63,24 +79,27 @@ suggesting, or modifying code within the 5. **Testing**: * Write unit tests using `EXPECT_EQ`, `EXPECT_TRUE`, etc. - * Use macros in `tsl/platform/status_matchers.h` instead of their TF_* - counterparts. For example: - * Use `ASSERT_OK_AND_ASSIGN`, `ASSERT_OK`, and `EXPECT_OK`. + * Use the status testing macros `ASSERT_OK_AND_ASSIGN`, `ASSERT_OK`, and + `EXPECT_OK` from `` and the status matchers + (`absl_testing::IsOk`, `absl_testing::IsOkAndHolds`, + `absl_testing::StatusIs`) from `absl/status/status_matchers.h` instead + of their legacy `TF_*` or `tsl::testing::*` counterparts: * DO NOT USE `TF_ASSERT_OK_AND_ASSIGN`, `TF_ASSERT_OK`, and `TF_EXPECT_OK`. - * When you refactor code that uses the TF_* macros., replace them, but - do not touch unrelated code. - * Put tests into an anonymous namespace + * When you refactor code that uses the `TF_*` macros, replace them, + but do not touch unrelated code. + * Put tests into an anonymous namespace. * Use `HloPjRtInterpreterReferenceMixin` or `HloHardwareIndependentTestBase` for compiler pass tests locally where possible. * Ensure tests are deterministic and do not flake. 6. **BUILD targets**: - * When defining BUILD targets prefer these XLA specific rules: + * When defining BUILD targets prefer these XLA specific rules: * Instead of `proto_library` use `tf_proto_library`. There is no need to define language specific targets with `tf_proto_library`. - * Instead of `cc_test` use `xla_cc_test`. + * Instead of `cc_test` use `xla_cc_test` (or `xla_test` when the test + requires a real hardware backend). 7. **Explicit Typing**: * **Avoid `auto`** in public headers or complex logic chains. @@ -92,8 +111,15 @@ suggesting, or modifying code within the * *Example*: Do not rely on layout information before Layout Assignment. 9. **Namespaces**: - * Prefer xla::gpu over nested namespaces. + * Prefer `xla::gpu` over nested namespaces. 10. **MLIR Operation Creation**: * **Always** use the static `OpTy::create(rewriter, ...)` method when creating MLIR operations. * **Avoid** using `rewriter.create(...)`. This syntax is deprecated. + +## Editing This File + +This file is shared between the open-source repository and Google-internal +development. Keep it in open-source terms only: repository-relative paths +(`xla/...`), OSS macro names (`ABSL_RETURN_IF_ERROR`, `ABSL_ASSIGN_OR_RETURN`), +Bazel labels (`//xla/...`, `@repo//...`), and `bazel` commands. 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 8329f81c5bceb4..88ec1af6050f4b 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 @@ -123,9 +123,11 @@ BlockLevelEmitterBackend::GetSupportedConfigs(const HloInstruction& instr) { ->config() .debug_options() .xla_gpu_fusion_autotune_top_k_configs(); - ABSL_ASSIGN_OR_RETURN(TopKTiledRunTimeDataOrError tiled_runtime_data, - indexing_performance_model_.TryFindTopKBestTilingsForFusion( - *fusion_adaptor, num_configs)); + ABSL_ASSIGN_OR_RETURN( + TopKTiledRunTimeDataOrError tiled_runtime_data, + indexing_performance_model_ + .TryFindTopKBestTilingsForFusionAsync(*fusion_adaptor, num_configs) + .Await()); if (std::holds_alternative(tiled_runtime_data)) { return std::vector>(); @@ -152,9 +154,10 @@ BlockLevelEmitterBackend::GetCostModelConfig(const HloInstruction& instr) { auto fusion_adaptor = HloFusionAdaptor::ForInstruction(Cast(&instr)); - ABSL_ASSIGN_OR_RETURN( - TiledRunTimeDataOrError tiled_runtime_data_or_error, - indexing_performance_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or_error, + indexing_performance_model_ + .TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* fusion_decision = std::get_if(&tiled_runtime_data_or_error)) { diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc index 0813a96f1e2904..d814f140920cd7 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc @@ -306,7 +306,8 @@ absl::StatusOr FissionBackend::EstimateFissionPrologueEpilogue( auto fusion_adaptor = HloFusionAdaptor::ForInstruction(instr); ABSL_ASSIGN_OR_RETURN( TiledRunTimeDataOrError tiled_data_or_decision, - indexing_cost_model.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* tiled_data = std::get_if(&tiled_data_or_decision)) { diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen.cc b/third_party/xla/xla/backends/gpu/autotuner/miopen.cc index a2e925c27c0d30..0b9e59b4f58a86 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/miopen.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/miopen.cc @@ -276,22 +276,15 @@ bool MIOpenBackend::IsSupported(const HloInstruction& instr) { absl::StatusOr> MIOpenBackend::GetDefaultConfig( const HloInstruction& instr) { - if (IsSupported(instr)) { - auto config = std::make_unique(); - config->mutable_algorithm()->set_algo_id(0); - return config; - } - return absl::InvalidArgumentError( - "MIOpen backend doesn't support getting a default config for this " - "instruction."); + return absl::UnimplementedError( + "MIOpen backend doesn't support getting a default config"); } absl::StatusOr>> -GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, - const HloModule* module, - se::StreamExecutor* stream_executor, - se::DeviceAddressAllocator* allocator, - se::Stream* stream) { +MIOpenBackend::GetConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + se::StreamExecutor* stream_executor, se::DeviceAddressAllocator* allocator, + se::Stream* stream) { if (stream_executor == nullptr) { return absl::InvalidArgumentError("Null stream executor is not supported."); } @@ -357,7 +350,7 @@ GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, gpu_conv_config.filter_descriptor, gpu_conv_params.filter_buf, gpu_conv_config.output_descriptor, gpu_conv_params.output_buf, gpu_conv_config.conv_desc, - /* use_fallback = */ false, &scratch_allocator, engine_options, + /* use_fallback = */ do_not_autotune_, &scratch_allocator, engine_options, &conv_runners)); std::vector> configs; @@ -373,10 +366,10 @@ GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, } absl::StatusOr>> -GetFusedConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr, - const HloModule* module, - se::StreamExecutor* stream_executor, - se::DeviceAddressAllocator* allocator) { +MIOpenBackend::GetFusedConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + se::StreamExecutor* stream_executor, + se::DeviceAddressAllocator* allocator) { if (stream_executor == nullptr) { return absl::InvalidArgumentError("Null stream executor is not supported."); } @@ -453,13 +446,6 @@ MIOpenBackend::GetSupportedConfigs(const HloInstruction& instr) { stream_executor(), allocator_); } - if (do_not_autotune_) { - ABSL_ASSIGN_OR_RETURN(auto default_config, GetDefaultConfig(instr)); - std::vector> configs; - configs.push_back(std::move(default_config)); - return std::move(configs); - } - return GetConvolutionCustomCallConfigs( custom_call_instr, custom_call_instr->GetModule(), stream_executor(), allocator_, /* stream */ nullptr); diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen.h b/third_party/xla/xla/backends/gpu/autotuner/miopen.h index aa42f86b25c55d..67d18c5c3ff132 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/miopen.h +++ b/third_party/xla/xla/backends/gpu/autotuner/miopen.h @@ -26,6 +26,8 @@ limitations under the License. #include "xla/backends/autotuner/codegen_backend.h" #include "xla/backends/gpu/autotuner/gpu_codegen_backend.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/ir/hlo_module.h" #include "xla/service/compiler.h" #include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/stream_executor.h" @@ -62,6 +64,19 @@ class MIOpenBackend : public GpuCodegenBackend { private: bool IsSupported(const HloInstruction& instr) override; + absl::StatusOr>> + GetConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + stream_executor::StreamExecutor* stream_executor, + stream_executor::DeviceAddressAllocator* allocator, + stream_executor::Stream* stream); + + absl::StatusOr>> + GetFusedConvolutionCustomCallConfigs( + const HloCustomCallInstruction* instr, const HloModule* module, + stream_executor::StreamExecutor* stream_executor, + stream_executor::DeviceAddressAllocator* allocator); + bool do_not_autotune_; stream_executor::DeviceAddressAllocator* allocator_; }; diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc index 500935b952009c..a7e1bc0ef63232 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc @@ -142,10 +142,7 @@ TEST_F(MIOpenBackendTest, GetDefaultConfigFromMIOpenCustomCall) { absl::StatusOr> config = backend_.GetDefaultConfig( (*hlo_module->entry_computation()->root_instruction()->operand(0))); - TF_ASSERT_OK(config); - ASSERT_TRUE((*config)->has_algorithm()); - MIOpenBackendConfig algorithm_config = (*config)->algorithm(); - EXPECT_EQ(algorithm_config.algo_id(), 0); + EXPECT_THAT(config, absl_testing::StatusIs(absl::StatusCode::kUnimplemented)); } TEST_F(MIOpenBackendTest, ApplyConfigToMIOpenCustomCall) { diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc index a40bcbd0aed1b3..6f45f3bc155372 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc @@ -47,7 +47,6 @@ limitations under the License. #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Location.h" @@ -1248,81 +1247,109 @@ absl::StatusOr CreateCollectiveKernelSpec( } } -CollectiveCodegenConfig CreateCollectiveCodegenConfig( - const HloInstruction* instr) { - const HloInstruction* collective = instr; - if (instr->opcode() == HloOpcode::kFusion) { - collective = instr->fused_instructions_computation()->root_instruction(); - } - CollectiveCodegenConfig config; - auto gpu_config = collective->backend_config(); - if (!gpu_config.ok()) { - return config; - } - const auto strategy = - gpu_config->collective_backend_config().kernel_strategy(); - // One-shot AllGather: the runtime copies input to symmetric scratch and the - // kernel needs a barrier before reading peers' data. - if (collective->opcode() == HloOpcode::kAllGather && - strategy == CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) { - config.copy_input_to_scratch = true; - config.emit_entry_barrier = true; - } - return config; -} +namespace { + +struct AllGatherRewriteContext { + mlir::stablehlo::AllGatherOp op; + xtile::EntryFuncOp entry_func; + xtile::ExtractTileOp input_extract; + xtile::SelectBufferOp pull_select; + mlir::Value rank_arg; + mlir::Value signal_buffers_arg; + int64_t world_size = 0; +}; -absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, - int32_t world_size) { - // Find the xtile::EntryFuncOp in the module. - xtile::EntryFuncOp entry_func = nullptr; - for (auto fn : module.getOps()) { - entry_func = fn; - break; +absl::StatusOr BuildAllGatherContext( + mlir::stablehlo::AllGatherOp op) { + if (op.getOperands().size() != 1) { + return absl::InvalidArgumentError( + "AllGather op must have exactly one operand."); } + auto entry_func = op->getParentOfType(); if (!entry_func) { - return absl::InternalError( - "No xtile::EntryFuncOp found in module for barrier insertion."); + return absl::InvalidArgumentError( + "AllGather op must be in an XTile entry function."); } - - // The opaque args are appended after the regular input/output args. - // Their count is stored in the "num_opaque_args" attribute. auto num_opaque_attr = entry_func->getAttrOfType("num_opaque_args"); if (!num_opaque_attr || num_opaque_attr.getInt() < kNumCollectiveMetadataArgs) { - return absl::InternalError( - absl::StrCat("Expected at least ", kNumCollectiveMetadataArgs, - " opaque args for collective entry barrier, got ", - num_opaque_attr ? num_opaque_attr.getInt() : 0)); + return absl::InvalidArgumentError( + "EntryFuncOp does not have collective metadata arguments."); } - // The opaque args start at (total_args - num_opaque_args). + mlir::Value input_tile = op.getOperand(0); + auto input_extract = llvm::dyn_cast_if_present( + input_tile.getDefiningOp()); + if (!input_extract && input_tile.getDefiningOp() && + input_tile.getDefiningOp()->getNumOperands() > 0) { + // Workaround(i1_to_i8_workaround): booleans are stored as i8 and cast to + // i1 after ExtractTileOp. + input_extract = llvm::dyn_cast_if_present( + input_tile.getDefiningOp()->getOperand(0).getDefiningOp()); + } + if (!input_extract) { + return absl::InvalidArgumentError( + "AllGather operand must be defined by xtile::ExtractTileOp."); + } + + auto pull_select = llvm::dyn_cast_if_present( + input_extract.getSource().getDefiningOp()); + if (!pull_select) { + return absl::InvalidArgumentError( + "AllGather ExtractTileOp source must be defined by " + "xtile::SelectBufferOp."); + } + + ABSL_ASSIGN_OR_RETURN(auto replica_groups, + xla::ConvertReplicaGroups(op.getReplicaGroups(), op)); + if (replica_groups->replica_groups().empty()) { + return absl::InvalidArgumentError( + "AllGather replica groups must not be empty."); + } + + AllGatherRewriteContext ctx; + ctx.op = op; + ctx.entry_func = entry_func; + ctx.input_extract = input_extract; + ctx.pull_select = pull_select; + ctx.world_size = replica_groups->num_devices_per_group(); + int32_t total_args = entry_func.getNumArguments(); - int32_t opaque_start = + int32_t metadata_args_start = total_args - num_opaque_attr.getInt() - kNumTileIndexArgs; // Layout: opaque[0]=rank, opaque[1]=invocation count (unused, the signal // value comes from a counter in device memory), opaque[2]=signal_buffers. - mlir::Value rank_arg = entry_func.getArgument(opaque_start); - mlir::Value signal_buffers_arg = entry_func.getArgument(opaque_start + 2); + ctx.rank_arg = entry_func.getArgument(metadata_args_start); + ctx.signal_buffers_arg = entry_func.getArgument(metadata_args_start + 2); + return ctx; +} - // Insert at the beginning of the entry block, right after program_id - // extraction (which is the first op). We want the barrier before any - // scf.for tile loop. - mlir::Block& entry_block = entry_func.front(); - auto loc = entry_func.getLoc(); - mlir::ImplicitLocOpBuilder builder(loc, &entry_block, entry_block.begin()); +} // namespace + +mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op, + mlir::PatternRewriter& rewriter) { + absl::StatusOr maybe_ctx = BuildAllGatherContext(op); + if (!maybe_ctx.ok()) { + return rewriter.notifyMatchFailure(op, maybe_ctx.status().message()); + } + AllGatherRewriteContext& ctx = *maybe_ctx; + + mlir::ImplicitLocOpBuilder builder(ctx.pull_select.getLoc(), rewriter); + builder.setInsertionPoint(ctx.pull_select); mlir::Value block_id = ttir::GetProgramIdOp::create(builder, 0); mlir::Value signal_value = EmitDeviceInvocationCount( - builder, signal_buffers_arg, rank_arg, block_id, world_size); + builder, ctx.signal_buffers_arg, ctx.rank_arg, block_id, ctx.world_size); // Inter-block barrier via signal flags. This blocks until all // remote ranks have also signaled. - mtx::BlockBarrierOp::create(builder, signal_buffers_arg, rank_arg, + mtx::BlockBarrierOp::create(builder, ctx.signal_buffers_arg, ctx.rank_arg, signal_value, /*signal_slot=*/nullptr, - builder.getI32IntegerAttr(world_size), + builder.getI32IntegerAttr(ctx.world_size), /*signal_stride=*/nullptr); - return absl::OkStatus(); + rewriter.replaceOp(op, op.getOperand(0)); + return mlir::success(); } } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h index 094e529a0ded1b..5088623ca49d13 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h @@ -24,7 +24,6 @@ limitations under the License. #include "absl/status/statusor.h" #include "llvm/ADT/SmallVector.h" #include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/Types.h" #include "mlir/Support/LLVM.h" @@ -101,23 +100,13 @@ absl::StatusOr> GetCollectiveUnmanagedKernelArguments( mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op, mlir::PatternRewriter& rewriter); -// Creates a lightweight codegen config from the collective HLO instruction. -// The returned config drives codegen decisions (e.g. whether the runtime must -// copy the input to scratch before kernel launch). -CollectiveCodegenConfig CreateCollectiveCodegenConfig( - const HloInstruction* instr); +// Rewrites stablehlo all-gather op to a triton implementation. +mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op, + mlir::PatternRewriter& rewriter); // Creates a CollectiveKernelSpec for a given collective or fusion instruction. absl::StatusOr CreateCollectiveKernelSpec( const HloInstruction* instr, const LaunchDimensions& launch_dimensions); -// Emits a collective entry barrier at the start of the entry function in -// |module|. The barrier ensures all ranks have completed their D2D copies -// before any rank starts reading from the symmetric scratch buffers. -// Uses the opaque metadata args (rank, signal_value, signal_buffers) already -// present in the EntryFuncOp. -absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, - int32_t world_size); - } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_CODEGEN_TRITON_COLLECTIVE_EMITTER_H_ diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir index bf680cada9aaeb..b4ed1827ab6ff3 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir @@ -80,7 +80,7 @@ func.func @insert_wrong_shape(%arg0: tensor<8xbf16>, %arg1: !tt.ptr) { func.func @block_barrier_stride_without_slot(%ptr: !tt.ptr, %rank: i32, %sig: i32) { // expected-error @+1 {{'triton_xla.block_barrier' op signal_stride requires signal_slot to be specified}} - triton_xla.block_barrier %ptr, %rank, %sig, {world_size = 2 : i32, signal_stride = 4 : i32} : (!tt.ptr, i32, i32) -> () + triton_xla.block_barrier %ptr, %rank, %sig : (!tt.ptr, i32, i32) -> () return } @@ -88,7 +88,7 @@ func.func @block_barrier_stride_without_slot(%ptr: !tt.ptr, %rank: i32, %si func.func @block_barrier_slot_without_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32, %slot: i32) { // expected-error @+1 {{'triton_xla.block_barrier' op signal_slot requires a positive signal_stride}} - triton_xla.block_barrier %ptr, %rank, %sig, %slot, {world_size = 2 : i32} : (!tt.ptr, i32, i32, i32) -> () + triton_xla.block_barrier %ptr, %rank, %sig, %slot : (!tt.ptr, i32, i32, i32) -> () return } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td index e978703b7ccdf5..d5c5e683ab6961 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_dialect.td @@ -28,6 +28,7 @@ def XlaTritonDialect : Dialect { let cppNamespace = "::mlir::triton::xla"; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 1; let dependentDialects = [ "mlir::triton::TritonDialect", "mlir::triton::gpu::TritonGPUDialect", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td index 27bbca78c5aad7..88aa3c2cf4a1e6 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td @@ -204,7 +204,7 @@ def TTXLA_SqueezeDimsOp : TTXLA_Op<"squeeze_dims", [ let summary = "Remove a size-1 dimension from a tensor."; let arguments = (ins TT_Tensor:$src, I32Attr:$axis); let results = (outs TT_Tensor:$result); - let assemblyFormat = "$src attr-dict `:` type($src) `->` type($result)"; + let assemblyFormat = "$src prop-dict attr-dict `:` type($src) `->` type($result)"; } def TTXLA_GetTidOp : TTXLA_Op<"get_tid", [Pure]> { @@ -411,7 +411,7 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { let results = (outs); let assemblyFormat = [{ - $signal_buffers `,` $device_rank `,` $signal_value `,` ($signal_slot^ `,`)? attr-dict `:` + $signal_buffers `,` $device_rank `,` $signal_value (`,` $signal_slot^)? prop-dict attr-dict `:` functional-type(operands, results) }]; @@ -449,7 +449,7 @@ def TTXLA_GetPeerPtrOp : TTXLA_Op<"get_peer_ptr", [Pure]> { let results = (outs Arg]>:$result); let assemblyFormat = [{ - $address `,` $peer_id `,` $metadata `,` attr-dict `:` + $address `,` $peer_id `,` $metadata prop-dict attr-dict `:` functional-type(operands, results) }]; } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD index 22cfd08b2cd3b6..1532db909a751e 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD @@ -68,7 +68,6 @@ xla_cc_binary( deviceless_tests = [ "collectives/all_reduce.hlo", "collectives/all_gather.hlo", - "collectives/all_gather_one_shot.hlo", "concat/concat_active_scalar_constant.hlo", "concat/concat_offset_not_a_multiple_of_tile_size.hlo", "constraints/max_tensor_elements_violation.hlo", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo index 39b3bc64e1b787..b065647800ae85 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo @@ -33,14 +33,27 @@ ENTRY entry { // CHECK: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<2xi64>, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<256x128xf32>, -// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: !tt.ptr, +// CHECK-SAME: %[[RANK:[a-zA-Z0-9_]+]]: i32, +// CHECK-SAME: %[[INVOCATION_COUNT:[a-zA-Z0-9_]+]]: i32, +// CHECK-SAME: %[[SIGNAL_BUFFERS:[a-zA-Z0-9_]+]]: !tt.ptr, // CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index {xla.range = [0 : index, 127 : index]}) -// CHECK: %[[IDX0:.*]] = xla.apply_indexing #[[MAP0]](%[[PID]]) -// CHECK: %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[PID]]) -// CHECK: %[[IDX2:.*]] = xla.apply_indexing #[[MAP2]](%[[PID]]) +// CHECK-DAG: %[[IDX0:.*]] = xla.apply_indexing #[[MAP0]](%[[PID]]) +// CHECK-DAG: %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[PID]]) +// CHECK-DAG: %[[IDX2:.*]] = xla.apply_indexing #[[MAP2]](%[[PID]]) +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 +// CHECK-DAG: %[[WORLD_SIZE:.*]] = arith.constant 2 : i32 +// CHECK: %[[BLOCK_ID:.*]] = tt.get_program_id x +// CHECK: %[[SIGNAL_BUFFER_ADDR:.*]] = tt.addptr %[[SIGNAL_BUFFERS]], %[[RANK]] +// CHECK: %[[SIGNAL_BUFFER_INT:.*]] = tt.load %[[SIGNAL_BUFFER_ADDR]] +// CHECK: %[[SIGNAL_BUFFER:.*]] = tt.int_to_ptr %[[SIGNAL_BUFFER_INT]] +// CHECK: %[[BLOCK_OFFSET:.*]] = arith.muli %[[BLOCK_ID]], %[[WORLD_SIZE]] +// CHECK: %[[COUNTER_IDX:.*]] = arith.addi %[[BLOCK_OFFSET]], %[[RANK]] +// CHECK: %[[COUNTER_PTR:.*]] = tt.addptr %[[SIGNAL_BUFFER]], %[[COUNTER_IDX]] +// CHECK: %[[COUNTER:.*]] = tt.load %[[COUNTER_PTR]] {isVolatile = true} +// CHECK: %[[SIGNAL_VALUE:.*]] = arith.addi %[[COUNTER]], %[[C1]] +// CHECK: triton_xla.block_barrier %[[SIGNAL_BUFFERS]], %[[RANK]], %[[SIGNAL_VALUE]] +// CHECK-SAME: // CHECK: %[[BUF:.*]] = xtile.select_buffer %[[ARG0]][%[[IDX2]]] : memref<2xi64> -> memref<128x128xf32> // CHECK: %[[EXT:.*]] = xtile.extract %[[BUF]][%[[IDX0]], %[[IDX1]]] [16, 16] [1, 1] : memref<128x128xf32> -> tensor<16x16xf32> // CHECK: %[[IDX3:.*]] = xla.apply_indexing #[[MAP3]](%[[PID]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo deleted file mode 100644 index f04d9c1b766fe6..00000000000000 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo +++ /dev/null @@ -1,58 +0,0 @@ -// RUN: fusion_to_triton %s | FileCheck %s -// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s -f { - %param0 = f32[128,128]{1,0} parameter(0) - ROOT result = f32[256,128]{1,0} all-gather(%param0), - replica_groups={{0,1}}, dimensions={0}, - backend_config={"collective_backend_config":{ - "kernel_strategy":"KERNEL_STRATEGY_TRITON_ONE_SHOT"}} -} - -ENTRY entry { - p0 = f32[128,128]{1,0} parameter(0) - ROOT %fusion = f32[256,128]{1,0} fusion(p0), - kind=kCustom, calls=f, - backend_config={ - "fusion_backend_config": { - "kind": "__triton_collective", - "block_level_fusion_config": { - "num_warps": "4", - "output_tiles": [{sizes: [16,16]}], - "num_ctas": 1, - "num_stages": 1, - "is_tma_allowed": false, - "is_warp_specialization_allowed": false - } - } - } -} - -// One-shot all-gather waits on an entry barrier before reading the peers' -// scratch buffers. The signal value of the barrier is derived from this block's -// own slot in the local signal buffer (`SignalBuffers[rank] -// [block_id * world_size + rank]`). - -// CHECK: xtile.entry_func @triton_fn( -// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<2xi64>, -// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<256x128xf32>, -// CHECK-SAME: %[[RANK:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[INVOCATION_COUNT:[a-zA-Z0-9_]+]]: i32, -// CHECK-SAME: %[[SIGNAL_BUFFERS:[a-zA-Z0-9_]+]]: !tt.ptr, -// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index{{( \{xla.range = \[0 : index, 127 : index\]\})?}}) -// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 -// CHECK-DAG: %[[WORLD_SIZE:.*]] = arith.constant 2 : i32 -// CHECK: %[[BLOCK_ID:.*]] = tt.get_program_id x -// CHECK: %[[SIGNAL_BUFFER_ADDR:.*]] = tt.addptr %[[SIGNAL_BUFFERS]], %[[RANK]] -// CHECK: %[[SIGNAL_BUFFER_INT:.*]] = tt.load %[[SIGNAL_BUFFER_ADDR]] -// CHECK: %[[SIGNAL_BUFFER:.*]] = tt.int_to_ptr %[[SIGNAL_BUFFER_INT]] -// CHECK: %[[BLOCK_OFFSET:.*]] = arith.muli %[[BLOCK_ID]], %[[WORLD_SIZE]] -// CHECK: %[[COUNTER_IDX:.*]] = arith.addi %[[BLOCK_OFFSET]], %[[RANK]] -// CHECK: %[[COUNTER_PTR:.*]] = tt.addptr %[[SIGNAL_BUFFER]], %[[COUNTER_IDX]] -// CHECK: %[[COUNTER:.*]] = tt.load %[[COUNTER_PTR]] {isVolatile = true} -// CHECK: %[[SIGNAL_VALUE:.*]] = arith.addi %[[COUNTER]], %[[C1]] -// CHECK: triton_xla.block_barrier %[[SIGNAL_BUFFERS]], %[[RANK]], %[[SIGNAL_VALUE]], -// CHECK-SAME: {world_size = 2 : i32} -// CHECK: xtile.select_buffer %[[ARG0]] -// CHECK: xtile.extract -// CHECK: xtile.insert -// CHECK: xtile.return diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo index 8ae71f7f0d3ba6..731791d3cbe9b0 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo @@ -43,8 +43,8 @@ fusion{ // CHECK-SAME: : tensor<4096x!tt.ptr> // CHECK: ttg.barrier local -// CHECK: triton_xla.block_barrier %[[ARG4]], %[[ARG2]], %[[SIGNAL_VALUE]], -// CHECK-SAME: {world_size = 2 : i32} : (!tt.ptr, i32, i32) +// CHECK: triton_xla.block_barrier %[[ARG4]], %[[ARG2]], %[[SIGNAL_VALUE]] +// CHECK-SAME: : (!tt.ptr, i32, i32) // CHECK: xtile.insert // CHECK-SAME: : tensor<4096xf32> -> memref<65536xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc index ce1a4f839e0350..5653ba9c4e6013 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc @@ -584,8 +584,10 @@ TEST_F(XTileDialectTest, HloAllGatherDotLowering) { CHECK: %[[SELECT2:.*]] = xtile.select_buffer %[[SELECT1]][%{{.*}}] CHECK-SAME: : memref<2xi64> -> memref<128x128xf32> CHECK: %[[LHS_TILE:.*]] = xtile.extract %[[SELECT2]] + CHECK: %[[AG1:.*]] = "stablehlo.all_gather"(%[[LHS_TILE]]) + CHECK: %[[AG2:.*]] = "stablehlo.all_gather"(%[[AG1]]) CHECK: %[[RHS_TILE:.*]] = xtile.extract %arg1 - CHECK: stablehlo.dot_general %[[LHS_TILE]], %[[RHS_TILE]] + CHECK: stablehlo.dot_general %[[AG2]], %[[RHS_TILE]] )")); } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc index c47c144a25d843..09747bc73d0261 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc @@ -1030,6 +1030,18 @@ class LowerDotGeneral : public mlir::OpRewritePattern { bool warp_specialization_allowed_; }; +class LowerAllGather : public mlir::OpRewritePattern { + public: + using OpRewritePattern::OpRewritePattern; + + private: + mlir::LogicalResult matchAndRewrite( + stablehlo::AllGatherOp op, + mlir::PatternRewriter& rewriter) const override { + return ::xla::gpu::RewriteAllGather(op, rewriter); + } +}; + class LowerAllReduce : public mlir::OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; @@ -1054,7 +1066,8 @@ class StableHLOLowerToTritonPass { mlir::RewritePatternSet patterns(mlir_context); patterns.add(mlir_context); + LowerReduce, LowerReshape, LowerAllGather, LowerAllReduce>( + mlir_context); patterns.add(mlir_context); patterns.add(mlir_context, warp_specialization_allowed_); if (mlir::failed(mlir::applyPatternsGreedily(getOperation(), diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir index 7fcbfbf52b4d08..94c56e69205642 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir @@ -45,7 +45,7 @@ tt.func @block_barrier_kernel( // CHECK-NEXT: } // CHECK-NEXT: ttg.barrier local // CHECK-NEXT: tt.return - triton_xla.block_barrier %ptr, %rank, %signal_value, { world_size = 8 : i32 } : + triton_xla.block_barrier %ptr, %rank, %signal_value : (!tt.ptr, i32, i32) -> () tt.return } @@ -88,7 +88,7 @@ tt.func @block_barrier_asymmetric_kernel( // CHECK-NEXT: } // CHECK-NEXT: ttg.barrier local // CHECK-NEXT: tt.return - triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot, { world_size = 2 : i32, signal_stride = 4 : i32 } : + triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot : (!tt.ptr, i32, i32, i32) -> () tt.return } \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir index ac41b44c900344..e9009950d7949d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_fold_transpose.mlir @@ -93,3 +93,59 @@ func.func @push_transpose_up_through_mask(%arg0: tensor<4x8xf32>, %arg1: f32) -> %1 = tt.trans %0 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> return %1 : tensor<8x4xf32> } + +// CHECK-LABEL: func @push_transpose_up_through_diamond_dag +func.func @push_transpose_up_through_diamond_dag( + %input: memref<4x8xf32, #xtile.layout<[1, 0]>>, + %offset0: index, %offset1: index) -> tensor<8x4xf32> { + // CHECK: %[[TRANSPOSE:.*]] = memref.transpose %arg0 (d0, d1) -> (d1, d0) + // CHECK-NEXT: %[[EXTRACT:.*]] = xtile.extract %[[TRANSPOSE]][%arg2, %arg1] [8, 4] [1, 1] + // CHECK-NEXT: %[[SQ:.*]] = arith.mulf %[[EXTRACT]], %[[EXTRACT]] : tensor<8x4xf32> + // CHECK-NEXT: %[[ADD:.*]] = arith.addf %[[EXTRACT]], %[[SQ]] : tensor<8x4xf32> + // CHECK-NEXT: %[[MUL:.*]] = arith.mulf %[[SQ]], %[[ADD]] : tensor<8x4xf32> + // CHECK-NOT: tt.trans + // CHECK-NEXT: return %[[MUL]] : tensor<8x4xf32> + %tile = xtile.extract %input[%offset0, %offset1][4, 8][1, 1] + : memref<4x8xf32, #xtile.layout<[1, 0]>> -> tensor<4x8xf32> + %sq = arith.mulf %tile, %tile : tensor<4x8xf32> + %add = arith.addf %tile, %sq : tensor<4x8xf32> + %mul = arith.mulf %sq, %add : tensor<4x8xf32> + %res = tt.trans %mul {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + return %res : tensor<8x4xf32> +} + +// CHECK-LABEL: func @push_transpose_up_through_if_multiple_uses +func.func @push_transpose_up_through_if_multiple_uses( + %arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>, %cond: i1) -> tensor<8x4xf32> { + // CHECK-DAG: %[[TRANS0:.*]] = tt.trans %arg0 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + // CHECK-DAG: %[[TRANS1:.*]] = tt.trans %arg1 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + // CHECK: %[[IF:.*]] = scf.if %arg2 -> (tensor<8x4xf32>) { + // CHECK: scf.yield %[[TRANS0]] : tensor<8x4xf32> + // CHECK: } else { + // CHECK: scf.yield %[[TRANS1]] : tensor<8x4xf32> + // CHECK: } + // CHECK-NEXT: %[[SQ:.*]] = arith.mulf %[[IF]], %[[IF]] : tensor<8x4xf32> + // CHECK-NEXT: %[[ADD:.*]] = arith.addf %[[IF]], %[[SQ]] : tensor<8x4xf32> + // CHECK-NOT: tt.trans + // CHECK-NEXT: return %[[ADD]] : tensor<8x4xf32> + %0 = scf.if %cond -> tensor<4x8xf32> { + scf.yield %arg0 : tensor<4x8xf32> + } else { + scf.yield %arg1 : tensor<4x8xf32> + } + %sq = arith.mulf %0, %0 : tensor<4x8xf32> + %add = arith.addf %0, %sq : tensor<4x8xf32> + %res = tt.trans %add {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + return %res : tensor<8x4xf32> +} + +// CHECK-LABEL: func @do_not_push_transpose_when_producer_has_non_transpose_user +func.func @do_not_push_transpose_when_producer_has_non_transpose_user( + %arg0: tensor<4x8xf32>) -> (tensor<4x8xf32>, tensor<8x4xf32>) { + // CHECK-NEXT: %[[NEG:.*]] = arith.negf %arg0 : tensor<4x8xf32> + // CHECK-NEXT: %[[TRANS:.*]] = tt.trans %[[NEG]] {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + // CHECK-NEXT: return %[[NEG]], %[[TRANS]] : tensor<4x8xf32>, tensor<8x4xf32> + %0 = arith.negf %arg0 : tensor<4x8xf32> + %1 = tt.trans %0 {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> + return %0, %1 : tensor<4x8xf32>, tensor<8x4xf32> +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir index 180108fe7f2b0a..313a6fd6265595 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_remote_access.mlir @@ -62,8 +62,8 @@ tt.func @get_peer_ptr( // Load metadata->buffer_root_ptrs[argument_offset + peer_id] + offset. // CHECK-NEXT: %11 = arith.addi %10, %6 : i64 // CHECK-NEXT: %12 = tt.int_to_ptr %11 : i64 -> !tt.ptr - %arg_0_peer_ptr = triton_xla.get_peer_ptr %arg0, %peer_id, %metadata, - { argument_index = 0 : i32, world_size = 2 : i32 } : + %arg_0_peer_ptr = triton_xla.get_peer_ptr %arg0, %peer_id, %metadata + : (!tt.ptr, i64, !tt.ptr) -> !tt.ptr // Load metadata->rank @@ -90,8 +90,8 @@ tt.func @get_peer_ptr( // CHECK-NEXT: %24 = arith.addi %23, %19 : i64 // CHECK-NEXT: %25 = tt.int_to_ptr %24 : i64 -> !tt.ptr - %arg_1_peer_ptr = triton_xla.get_peer_ptr %arg1, %peer_id, %metadata, - { argument_index = 1 : i32, world_size = 2 : i32 } : + %arg_1_peer_ptr = triton_xla.get_peer_ptr %arg1, %peer_id, %metadata + : (!tt.ptr, i64, !tt.ptr) -> !tt.ptr // Avoid optimizing away the get_peer_ptr calls, by returning xor of the two diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir index 63d0cbc0faa245..4e87ba0dc1048d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_squeeze_dims.mlir @@ -22,7 +22,7 @@ func.func @push_squeeze_dims_up_through_elementwise(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // CHECK: arith.negf {{.*}} : tensor<4x8xf32> %0 = arith.negf %arg0 : tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -32,7 +32,7 @@ func.func @push_squeeze_dims_up_through_elementwise(%arg0: tensor<4x1x8xf32>) -> func.func @push_squeeze_dims_up_through_multiple_results(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // CHECK: tt.elementwise_inline_asm {{.*}} : tensor<4x8xf32> -> tensor<4x8xf32>, tensor<4x8xf32> %0:2 = tt.elementwise_inline_asm "" {constraints = "", packed_element = 1 : i32, pure = true} %arg0 : tensor<4x1x8xf32> -> tensor<4x1x8xf32>, tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0#0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0#0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -42,7 +42,7 @@ func.func @push_squeeze_dims_up_through_multiple_results(%arg0: tensor<4x1x8xf32 func.func @push_squeeze_dims_up_through_broadcast(%arg0: tensor<1x4x1x8xf32>) -> tensor<4x16x8xf32> { // CHECK: tt.broadcast {{.*}} : tensor<4x1x8xf32> -> tensor<4x16x8xf32> %0 = tt.broadcast %arg0 : tensor<1x4x1x8xf32> -> tensor<1x4x16x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 0 : i32} : tensor<1x4x16x8xf32> -> tensor<4x16x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<1x4x16x8xf32> -> tensor<4x16x8xf32> return %1 : tensor<4x16x8xf32> } @@ -52,7 +52,7 @@ func.func @push_squeeze_dims_up_through_broadcast(%arg0: tensor<1x4x1x8xf32>) -> func.func @push_squeeze_dims_up_through_trans(%arg0: tensor<4x1x8xf32>) -> tensor<8x4xf32> { // CHECK: tt.trans {{.*}} {order = array} : tensor<4x8xf32> -> tensor<8x4xf32> %0 = tt.trans %arg0 {order = array} : tensor<4x1x8xf32> -> tensor<8x4x1xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 2 : i32} : tensor<8x4x1xf32> -> tensor<8x4xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<8x4x1xf32> -> tensor<8x4xf32> return %1 : tensor<8x4xf32> } @@ -62,7 +62,7 @@ func.func @push_squeeze_dims_up_through_trans(%arg0: tensor<4x1x8xf32>) -> tenso func.func @push_squeeze_dims_up_through_join(%arg0: tensor<1x4xf32>, %arg1: tensor<1x4xf32>) -> tensor<4x2xf32> { // CHECK-DAG: tt.join {{.*}} : tensor<4xf32> -> tensor<4x2xf32> %0 = tt.join %arg0, %arg1 : tensor<1x4xf32> -> tensor<1x4x2xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 0 : i32} : tensor<1x4x2xf32> -> tensor<4x2xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<1x4x2xf32> -> tensor<4x2xf32> return %1 : tensor<4x2xf32> } @@ -77,7 +77,7 @@ func.func @push_squeeze_dims_up_through_reduce(%arg0: tensor<8x4x1xf32>) -> tens tt.reduce.return %1 : f32 // CHECK: }) : (tensor<8x4xf32>) -> tensor<8xf32> }) : (tensor<8x4x1xf32>) -> tensor<8x1xf32> - %2 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<8x1xf32> -> tensor<8xf32> + %2 = triton_xla.squeeze_dims %0 : tensor<8x1xf32> -> tensor<8xf32> return %2 : tensor<8xf32> } @@ -88,7 +88,7 @@ func.func @fold_squeeze_of_expand_cancelling(%arg0: tensor<4x8xf32>) -> tensor<4 // CHECK-NOT: tt.expand_dims // CHECK-NOT: triton_xla.squeeze_dims %0 = tt.expand_dims %arg0 {axis = 1 : i32} : tensor<4x8xf32> -> tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -96,10 +96,10 @@ func.func @fold_squeeze_of_expand_cancelling(%arg0: tensor<4x8xf32>) -> tensor<4 // CHECK-LABEL: func @fold_squeeze_of_expand_swapping func.func @fold_squeeze_of_expand_swapping(%arg0: tensor<4x1x8xf32>) -> tensor<1x4x8xf32> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1x8xf32> -> tensor<4x8xf32> // CHECK: tt.expand_dims {{.*}} {axis = 0 : i32} : tensor<4x8xf32> -> tensor<1x4x8xf32> %0 = tt.expand_dims %arg0 {axis = 0 : i32} : tensor<4x1x8xf32> -> tensor<1x4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 2 : i32} : tensor<1x4x1x8xf32> -> tensor<1x4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<1x4x1x8xf32> -> tensor<1x4x8xf32> return %1 : tensor<1x4x8xf32> } @@ -107,8 +107,8 @@ func.func @fold_squeeze_of_expand_swapping(%arg0: tensor<4x1x8xf32>) -> tensor<1 // CHECK-LABEL: func @squeeze_reshape func.func @squeeze_reshape(%arg0: tensor<4x1x1xf32>) -> tensor<4xf32> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1x1xf32> -> tensor<4x1xf32> - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1xf32> -> tensor<4xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1x1xf32> -> tensor<4x1xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1xf32> -> tensor<4xf32> %0 = tt.reshape %arg0 : tensor<4x1x1xf32> -> tensor<4xf32> return %0 : tensor<4xf32> } @@ -152,7 +152,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} { // CHECK-LABEL: func @reshape_with_encoding // CHECK-SAME: tensor<1x32xf32, #[[ARG_ENC:.+]]>) -> tensor<32xf32, #[[RES_ENC:.+]]> func.func @reshape_with_encoding(%arg0: tensor<1x32xf32, #arg_enc>) -> tensor<32xf32, #res_enc> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 0 : i32} : + // CHECK: triton_xla.squeeze_dims {{.*}} : // CHECK-SAME: tensor<1x32xf32, #[[ARG_ENC]]> -> tensor<32xf32, #[[RES_ENC]]> %0 = tt.reshape %arg0 : tensor<1x32xf32, #arg_enc> -> tensor<32xf32, #res_enc> return %0 : tensor<32xf32, #res_enc> @@ -171,7 +171,7 @@ func.func @fold_squeeze_dims_of_extract( %tile = xtile.extract %input[%offset, %offset, %offset][4, 1, 8][1, 1, 1] : memref<4x16x8xf32, #xtile.layout<[2, 1, 0]>> -> tensor<4x1x8xf32> // CHECK-NOT: triton_xla.squeeze_dims - %squeezed = triton_xla.squeeze_dims %tile {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %squeezed = triton_xla.squeeze_dims %tile : tensor<4x1x8xf32> -> tensor<4x8xf32> // CHECK: return %[[EXTRACT]] return %squeezed : tensor<4x8xf32> } @@ -187,7 +187,7 @@ func.func @squeeze_insert( %arg1: tensor<4x1x8xf32>, %offset: index) { // CHECK: %[[REDUCED:.*]] = triton_xla.squeeze_dims %[[TILE]] - // CHECK-SAME: {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + // CHECK-SAME: : tensor<4x1x8xf32> -> tensor<4x8xf32> // CHECK: xtile.insert %[[REDUCED]] into %[[BUFFER]] // CHECK-SAME: tensor<4x8xf32> -> memref<4x16x8xf32, #xtile.layout<[2, 1, 0]>> xtile.insert %arg1 into %arg0[%offset, %offset, %offset][4, 1, 8][1, 1, 1] @@ -213,10 +213,10 @@ func.func @squeeze_insert_unit_tensor( // CHECK-LABEL: func @reorder_squeeze_dims func.func @reorder_squeeze_dims(%arg0: tensor<4x1x8x1xf32>) -> tensor<4x8xf32> { - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 1 : i32} : tensor<4x1x8x1xf32> -> tensor<4x8x1xf32> - // CHECK: triton_xla.squeeze_dims {{.*}} {axis = 2 : i32} : tensor<4x8x1xf32> -> tensor<4x8xf32> - %0 = triton_xla.squeeze_dims %arg0 {axis = 3 : i32} : tensor<4x1x8x1xf32> -> tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x1x8x1xf32> -> tensor<4x8x1xf32> + // CHECK: triton_xla.squeeze_dims {{.*}} : tensor<4x8x1xf32> -> tensor<4x8xf32> + %0 = triton_xla.squeeze_dims %arg0 : tensor<4x1x8x1xf32> -> tensor<4x1x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } @@ -227,7 +227,7 @@ func.func @diamond(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // CHECK-NOT: arith.negf {{.*}} : tensor<4x1x8xf32> %0 = arith.negf %arg0 : tensor<4x1x8xf32> %1 = arith.addf %0, %0 : tensor<4x1x8xf32> - %2 = triton_xla.squeeze_dims %1 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %2 = triton_xla.squeeze_dims %1 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %2 : tensor<4x8xf32> } @@ -246,7 +246,7 @@ func.func @insert_expand_dims(%arg0: tensor<4x1xf32>) -> tensor<4xf32> { }) : (tensor<4x1xf32>) -> tensor<4xf32> %3 = tt.expand_dims %1 {axis = 1 : i32} : tensor<4xf32> -> tensor<4x1xf32> %4 = arith.addf %0, %3 : tensor<4x1xf32> - %5 = triton_xla.squeeze_dims %4 {axis = 1 : i32} : tensor<4x1xf32> -> tensor<4xf32> + %5 = triton_xla.squeeze_dims %4 : tensor<4x1xf32> -> tensor<4xf32> return %5 : tensor<4xf32> } @@ -267,7 +267,7 @@ func.func @push_squeeze_dims_up_through_if( scf.yield %arg0, %arg1 : tensor<16xf32>, tensor<4x1xf32> } // CHECK-NOT: triton_xla.squeeze_dims - %squeeze = triton_xla.squeeze_dims %if#1 {axis = 1 : i32} + %squeeze = triton_xla.squeeze_dims %if#1 : tensor<4x1xf32> -> tensor<4xf32> // CHECK: return return %if#0, %squeeze : tensor<16xf32>, tensor<4xf32> @@ -278,7 +278,7 @@ func.func @push_squeeze_dims_up_through_if( // FINALIZE-LABEL: func @squeeze_dims_to_reshape func.func @squeeze_dims_to_reshape(%arg0: tensor<4x1x8xf32>) -> tensor<4x8xf32> { // FINALIZE: tt.reshape {{.*}} : tensor<4x1x8xf32> -> tensor<4x8xf32> - %0 = triton_xla.squeeze_dims %arg0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %0 = triton_xla.squeeze_dims %arg0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %0 : tensor<4x8xf32> } @@ -289,6 +289,6 @@ func.func @push_squeeze_dims_up_through_mask( %arg0: tensor<4x1x8xf32>, %arg1: f32) -> tensor<4x8xf32> { // CHECK: xtile.mask %{{.*}} bounds [4, 6], %arg1 : tensor<4x8xf32> %0 = xtile.mask %arg0 bounds [4, 1, 6], %arg1 : tensor<4x1x8xf32> - %1 = triton_xla.squeeze_dims %0 {axis = 1 : i32} : tensor<4x1x8xf32> -> tensor<4x8xf32> + %1 = triton_xla.squeeze_dims %0 : tensor<4x1x8xf32> -> tensor<4x8xf32> return %1 : tensor<4x8xf32> } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc index 80d77ef1fcf3cc..790918f99f2f31 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_fold_transpose_pass.cc @@ -20,6 +20,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/ErrorHandling.h" @@ -57,6 +58,24 @@ namespace { return guard; } +bool CanFoldTransposeThroughProducer(TransOp op) { + return absl::c_all_of(op.getSrc().getUsers(), [&](Operation* user) { + auto other_trans = dyn_cast(user); + return other_trans && other_trans.getOrder() == op.getOrder(); + }); +} + +void ReplaceAllTransposeUsersWith(PatternRewriter& rewriter, Value old_src, + Value new_src) { + for (Operation* user : llvm::make_early_inc_range(old_src.getUsers())) { + rewriter.replaceOp(user, new_src); + } + if (Operation* producer = old_src.getDefiningOp(); + producer && producer->use_empty()) { + rewriter.eraseOp(producer); + } +} + // Push the transpose up through the extract tile, this will then be folded into // MemrefToPtr at the lowering stage. LogicalResult PushTransposeThroughExtractTile(TransOp op, @@ -65,6 +84,10 @@ LogicalResult PushTransposeThroughExtractTile(TransOp op, if (!extract) { return rewriter.notifyMatchFailure(op, "Transpose source is not extract."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } SmallVector reduced_dims = to_vector(extract.getReducedDimensions()); @@ -107,19 +130,18 @@ LogicalResult PushTransposeThroughExtractTile(TransOp op, return result; }; + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, extract); auto permutation_map = mlir::AffineMapAttr::get( mlir::AffineMap::getPermutationMap(permutation, rewriter.getContext())); // TODO(willfroom): Return a permutation layout (b/455478641). auto pushed_transpose = mlir::memref::TransposeOp::create( rewriter, extract.getLoc(), extract.getSource(), permutation_map); - rewriter.replaceOpWithNewOp<::xla::xtile::ExtractTileOp>( - op, op.getType(), pushed_transpose, permute(extract.getOffsets()), - permute(extract.getFullTileShape()), permute(extract.getStrides())); - - if (extract->use_empty()) { - rewriter.eraseOp(extract); - } + Value new_extract = ::xla::xtile::ExtractTileOp::create( + rewriter, extract.getLoc(), op.getType(), pushed_transpose, + permute(extract.getOffsets()), permute(extract.getFullTileShape()), + permute(extract.getStrides())); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_extract); return success(); } @@ -131,9 +153,16 @@ LogicalResult PushTransposeUpThroughBroadcast(TransOp op, return rewriter.notifyMatchFailure( // op, "Transpose source is not a broadcast."); } - Value new_trans = TransOp::create(rewriter, op.getLoc(), broadcast.getSrc(), - op.getOrderAttr()); - rewriter.replaceOpWithNewOp(op, op.getType(), new_trans); + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, broadcast); + Value new_trans = TransOp::create(rewriter, broadcast.getLoc(), + broadcast.getSrc(), op.getOrderAttr()); + Value new_broadcast = BroadcastOp::create(rewriter, broadcast.getLoc(), + op.getType(), new_trans); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_broadcast); return success(); } @@ -144,6 +173,10 @@ LogicalResult PushTransposeUpThroughExpandDims(TransOp op, return rewriter.notifyMatchFailure( op, "Transpose source is not an expand_dims."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } unsigned new_axis = [&] { for (auto [i, dim] : llvm::enumerate(op.getOrder())) { @@ -160,10 +193,12 @@ LogicalResult PushTransposeUpThroughExpandDims(TransOp op, dim -= dim > expand_dims.getAxis(); } - Value new_trans = - TransOp::create(rewriter, op.getLoc(), expand_dims.getSrc(), new_order); - rewriter.replaceOpWithNewOp(op, op.getType(), new_trans, - new_axis); + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, expand_dims); + Value new_trans = TransOp::create(rewriter, expand_dims.getLoc(), + expand_dims.getSrc(), new_order); + Value new_expand_dims = ExpandDimsOp::create( + rewriter, expand_dims.getLoc(), op.getType(), new_trans, new_axis); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_expand_dims); return success(); } @@ -175,13 +210,23 @@ LogicalResult PushTransposeUpThroughElementwise(TransOp op, return rewriter.notifyMatchFailure( op, "source is not a single-result elementwise op"); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, elementwise); + llvm::SmallDenseMap transposed_operands; SmallVector new_operands; new_operands.reserve(elementwise->getNumOperands()); for (Value operand : elementwise->getOperands()) { - if (auto tensor_type = dyn_cast(operand.getType())) { - operand = TransOp::create(rewriter, elementwise->getLoc(), operand, - op.getOrderAttr()); + if (isa(operand.getType())) { + auto [it, inserted] = transposed_operands.try_emplace(operand, nullptr); + if (inserted) { + it->second = TransOp::create(rewriter, elementwise->getLoc(), operand, + op.getOrderAttr()); + } + operand = it->second; } new_operands.push_back(operand); } @@ -189,7 +234,7 @@ LogicalResult PushTransposeUpThroughElementwise(TransOp op, Operation* new_op = rewriter.clone(*elementwise); new_op->setOperands(new_operands); new_op->getResult(0).setType(op.getType()); - rewriter.replaceOp(op, new_op->getResults()); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_op->getResult(0)); return success(); } @@ -213,19 +258,24 @@ LogicalResult PushTransposeUpThroughElementwise(TransOp op, LogicalResult PushTransposeUpIntoIf(TransOp op, PatternRewriter& rewriter) { Value src = op.getSrc(); auto if_op = src.getDefiningOp(); - if (!if_op || !src.hasOneUse()) { + if (!if_op) { return rewriter.notifyMatchFailure(op, "Expected scf.if producer."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } // Compute the new types for the if op. unsigned result_number = cast(src).getResultNumber(); auto new_types = llvm::to_vector(if_op.getResultTypes()); new_types[result_number] = op.getType(); - auto new_if_op = - scf::IfOp::create(rewriter, op.getLoc(), new_types, if_op.getCondition(), - /*addThenBlock=*/false, - /*addElseBlock=*/false); + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, if_op); + auto new_if_op = scf::IfOp::create(rewriter, if_op.getLoc(), new_types, + if_op.getCondition(), + /*addThenBlock=*/false, + /*addElseBlock=*/false); // Update then and else regions. for (auto [old_region, new_region] : @@ -235,13 +285,16 @@ LogicalResult PushTransposeUpIntoIf(TransOp op, PatternRewriter& rewriter) { continue; } auto yield_op = new_region->front().getTerminator(); - OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, yield_op); + OpBuilder::InsertionGuard yield_guard = + SetInsertionPoint(rewriter, yield_op); auto trans_op = - TransOp::create(rewriter, op.getLoc(), op.getType(), + TransOp::create(rewriter, if_op.getLoc(), op.getType(), yield_op->getOperand(result_number), op.getOrderAttr()); yield_op->setOperand(result_number, trans_op); } - rewriter.replaceOp(op, new_if_op.getResult(result_number)); + for (Operation* user : llvm::make_early_inc_range(src.getUsers())) { + rewriter.replaceOp(user, new_if_op.getResult(result_number)); + } rewriter.replaceOp(if_op, new_if_op); return success(); } @@ -255,7 +308,7 @@ LogicalResult HoistTransposeUpFromIf(TransOp op, PatternRewriter& rewriter) { return rewriter.notifyMatchFailure(op, "Operand defined inside scf.if."); } - op->moveBefore(if_op); + rewriter.modifyOpInPlace(op, [&] { op->moveBefore(if_op); }); return success(); } @@ -274,6 +327,10 @@ LogicalResult PushTransposeUpThroughReshape(TransOp op, return rewriter.notifyMatchFailure(op, "Transpose source is not a reshape."); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } auto operand_shape = reshape.getSrc().getType().getShape(); auto reshape_shape = reshape.getType().getShape(); @@ -311,9 +368,12 @@ LogicalResult PushTransposeUpThroughReshape(TransOp op, } } + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, reshape); auto new_trans = TransOp::create(rewriter, reshape.getLoc(), reshape.getSrc(), new_order); - rewriter.replaceOpWithNewOp(op, op.getType(), new_trans); + Value new_reshape = + ReshapeOp::create(rewriter, reshape.getLoc(), op.getType(), new_trans); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_reshape); return success(); } @@ -323,17 +383,24 @@ LogicalResult PushTransposeUpThroughMask(TransOp op, if (!mask_op) { return rewriter.notifyMatchFailure(op, "source is not a mask op"); } + if (!CanFoldTransposeThroughProducer(op)) { + return rewriter.notifyMatchFailure( + op, "Transpose source has users that are not matching transposes."); + } llvm::SmallVector new_bounds(op.getOrder().size()); for (auto [idx, dim] : llvm::enumerate(op.getOrder())) { new_bounds[idx] = mask_op.getBounds()[dim]; } - auto new_transpose = TransOp::create(rewriter, op.getLoc(), + OpBuilder::InsertionGuard guard = SetInsertionPoint(rewriter, mask_op); + auto new_transpose = TransOp::create(rewriter, mask_op.getLoc(), mask_op.getSource(), op.getOrderAttr()); - rewriter.replaceOpWithNewOp<::xla::xtile::MaskOp>( - op, op.getType(), new_transpose, new_bounds, mask_op.getValue()); + Value new_mask = ::xla::xtile::MaskOp::create(rewriter, mask_op.getLoc(), + op.getType(), new_transpose, + new_bounds, mask_op.getValue()); + ReplaceAllTransposeUsersWith(rewriter, op.getSrc(), new_mask); return success(); } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc index d52453328a497e..7cf600930a5697 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc @@ -447,14 +447,6 @@ absl::StatusOr CreateTritonModule( absl::MakeSpan(opaque_args_types), mlir_context, use_experimental_tiling, enable_same_shape_multi_output_fusion)); - if (fusion_kind == kTritonCollectiveFusionKind && - CreateCollectiveCodegenConfig(&fusion).emit_entry_barrier) { - const HloInstruction* root = hlo_computation->root_instruction(); - int32_t world_size = root->replica_groups()[0].replica_ids_size(); - ABSL_RETURN_IF_ERROR( - EmitCollectiveEntryBarrier(triton_module.get(), world_size)); - } - if (DumpingEnabledForHloModule(*hlo_computation->parent()) && DumpingEnabledForEmitter("triton-fusion", debug_options)) { auto suffix = absl::StrCat(fusion.name(), ".before_validation.ttir.txt"); diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 8fcdd36e3b9583..f5c3acd8945ea8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -5433,6 +5433,64 @@ xla_cc_test( ], ) +cc_library( + name = "device_slot", + hdrs = ["device_slot.h"], + visibility = ["//visibility:private"], + deps = [ + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/functional:function_ref", + ], +) + +cc_library( + name = "vector_storage", + srcs = ["vector_storage.cc"], + hdrs = ["vector_storage.h"], + visibility = ["//visibility:private"], + deps = [":device_slot"], +) + +xla_cc_test( + name = "vector_storage_test", + srcs = ["vector_storage_test.cc"], + deps = [ + ":device_slot", + ":vector_storage", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "cow_storage", + srcs = ["cow_storage.cc"], + hdrs = ["cow_storage.h"], + visibility = ["//visibility:private"], + deps = [ + ":device_slot", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:inlined_vector", + "@com_google_absl//absl/synchronization", + ], +) + +xla_cc_test( + name = "cow_storage_test", + srcs = ["cow_storage_test.cc"], + deps = [ + ":cow_storage", + ":device_slot", + "//xla/tsl/platform:env", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "custom_kernel_thunk", srcs = ["custom_kernel_thunk.cc"], diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc index 7ac59489518d7a..f0e781ec3bb726 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc @@ -232,7 +232,6 @@ absl::StatusOr CreateAllGatherKernelSpec( CollectiveKernelSpec kernel_spec = { /* .codegen_config= */ { /* .copy_input_to_scratch= */ true, - /* .emit_entry_barrier= */ true, /* .input_buffer_specs= */ {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, /* .output_buffer_specs= */ diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc index 8887f57cea1780..627337b25436bc 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc @@ -466,7 +466,6 @@ absl::StatusOr CreateAllReduceKernelSpec( CollectiveKernelSpec kernel_spec = { /* .codegen_config= */ { /* .copy_input_to_scratch= */ false, - /* .emit_entry_barrier= */ false, /* .input_buffer_specs= */ {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, /* .output_buffer_specs= */ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc index 79192758a8ddb9..03b534f1da4d8f 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc @@ -203,7 +203,6 @@ CollectiveKernelSpec CreateCollectiveKernelSpec( return { /*codegen_config=*/{ /*copy_input_to_scratch=*/false, - /*emit_entry_barrier=*/false, /*input_buffer_specs=*/ {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, /*output_buffer_specs=*/ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_params.h b/third_party/xla/xla/backends/gpu/runtime/collective_params.h index 2c2e97224e1f13..99fc121d4319b8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_params.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_params.h @@ -165,9 +165,6 @@ struct CollectiveCodegenConfig { // buffer before kernel launch. The kernel receives the scratch buffer as // its input argument. bool copy_input_to_scratch = false; - // If true, a cross-rank barrier is emitted before the tile loop. The barrier - // waits until all ranks have populated the symmetric scratch buffers. - bool emit_entry_barrier = false; // Specs for input operand buffers. std::vector input_buffer_specs; // Specs for output result buffers. diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc index fbc9a936b4d288..35618e2b64f268 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc @@ -204,6 +204,8 @@ std::optional GetCommandBufferCmdType( case Thunk::kPartitionId: case Thunk::kReplicaId: return DebugOptions::FUSION; + case Thunk::kCollectiveKernel: + return DebugOptions::COLLECTIVES_KERNEL; case Thunk::kWhile: return DebugOptions::WHILE; case Thunk::kConditional: @@ -214,7 +216,6 @@ std::optional GetCommandBufferCmdType( case Thunk::kAllReduce: case Thunk::kAllToAll: case Thunk::kCollectiveBroadcast: - case Thunk::kCollectiveKernel: case Thunk::kCollectivePermute: case Thunk::kCollectiveReduce: case Thunk::kGroup: diff --git a/third_party/xla/xla/backends/gpu/runtime/cow_storage.cc b/third_party/xla/xla/backends/gpu/runtime/cow_storage.cc new file mode 100644 index 00000000000000..9b86fc78b28fcb --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/cow_storage.cc @@ -0,0 +1,78 @@ +/* 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/backends/gpu/runtime/cow_storage.h" + +#include +#include +#include +#include + +#include "absl/container/inlined_vector.h" +#include "absl/synchronization/mutex.h" +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +DeviceSlot* CowStorage::Find(int device_ordinal) const { + if (device_ordinal < 0) { + return nullptr; + } + const Snapshot* snapshot = snapshot_.load(std::memory_order_acquire); + if (snapshot == nullptr) { + return nullptr; + } + return FindIn(*snapshot, device_ordinal); +} + +DeviceSlot* CowStorage::GetOrCreate(int device_ordinal, + DeviceSlotFactoryRef factory) { + if (device_ordinal < 0) { + return nullptr; + } + if (DeviceSlot* slot = Find(device_ordinal)) { + return slot; + } + absl::MutexLock lock(mutex_); + // Re-check under `mutex_` in case another thread concurrently created the + // slot for `device_ordinal` after the lock-free `Find` above. Snapshots are + // only published under `mutex_`, so relaxed is enough. + const Snapshot* current = snapshot_.load(std::memory_order_relaxed); + if (current != nullptr) { + if (DeviceSlot* slot = FindIn(*current, device_ordinal)) { + return slot; + } + } + slots_.push_back(factory()); + std::unique_ptr next = current == nullptr + ? std::make_unique() + : std::make_unique(*current); + next->push_back(Entry{device_ordinal, slots_.back().get()}); + snapshots_.push_back(std::move(next)); + const Snapshot* published = snapshots_.back().get(); + snapshot_.store(published, std::memory_order_release); + return published->back().slot; +} + +DeviceSlot* CowStorage::FindIn(const Snapshot& snapshot, int device_ordinal) { + for (const Entry& entry : snapshot) { + if (entry.device_ordinal == device_ordinal) { + return entry.slot; + } + } + return nullptr; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/cow_storage.h b/third_party/xla/xla/backends/gpu/runtime/cow_storage.h new file mode 100644 index 00000000000000..d7eae435c08554 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/cow_storage.h @@ -0,0 +1,80 @@ +/* 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_BACKENDS_GPU_RUNTIME_COW_STORAGE_H_ +#define XLA_BACKENDS_GPU_RUNTIME_COW_STORAGE_H_ + +#include +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/inlined_vector.h" +#include "absl/synchronization/mutex.h" +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +// Copy-on-write storage for holding per-device state slots built on first +// touch. The storage itself is thread-safe for concurrent lookups (`Find` is +// lock-free) and insertions (`GetOrCreate` serializes via a mutex and publishes +// a new immutable snapshot). However, the returned state slots are not +// internally synchronized; it is only safe to mutate them across threads if +// each thread is dedicated to a distinct device ordinal/index. +class CowStorage { + public: + CowStorage() = default; + + CowStorage(const CowStorage&) = delete; + CowStorage& operator=(const CowStorage&) = delete; + + // Lock-free. Returns nullptr if `device_ordinal` is negative or has no slot + // yet. + DeviceSlot* Find(int device_ordinal) const; + + // Returns the slot for `device_ordinal`, building it on the first call. + // Lock-free once the slot exists. Returns nullptr if `device_ordinal` is + // negative. + DeviceSlot* GetOrCreate(int device_ordinal, DeviceSlotFactoryRef factory) + ABSL_LOCKS_EXCLUDED(mutex_); + + private: + struct Entry { + int device_ordinal; + DeviceSlot* slot; + }; + + // Sized for a typical 8-GPU host so snapshots stay inline in the common + // case. + static constexpr size_t kInlinedSnapshotSize = 8; + + // Immutable once published. + using Snapshot = absl::InlinedVector; + + static DeviceSlot* FindIn(const Snapshot& snapshot, int device_ordinal); + + std::atomic snapshot_{nullptr}; + absl::Mutex mutex_; + // Heap slots keep their address across snapshot copies and stay cache-line + // aligned. + std::vector> slots_ ABSL_GUARDED_BY(mutex_); + // Freed only in the destructor, so readers never see a freed snapshot. + std::vector> snapshots_ ABSL_GUARDED_BY(mutex_); +}; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_RUNTIME_COW_STORAGE_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc b/third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc new file mode 100644 index 00000000000000..f5e24d835914f0 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/cow_storage_test.cc @@ -0,0 +1,304 @@ +/* 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/backends/gpu/runtime/cow_storage.h" + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/base/optimization.h" +#include "absl/container/flat_hash_set.h" +#include "absl/synchronization/notification.h" +#include "absl/time/time.h" +#include "xla/backends/gpu/runtime/device_slot.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/threadpool.h" + +namespace xla::gpu { +namespace { + +using ::testing::Contains; +using ::testing::Each; +using ::testing::Not; + +// Counts constructions and destructions. Neither copyable nor movable, so a +// storage that copies or moves its states fails to compile. +struct Tracked { + static constexpr int kInitialValue = 42; + + Tracked() { constructed.fetch_add(1); } + ~Tracked() { destroyed.fetch_add(1); } + Tracked(const Tracked&) = delete; + Tracked& operator=(const Tracked&) = delete; + + static void Reset() { + constructed.store(0); + destroyed.store(0); + } + + static inline std::atomic constructed{0}; + static inline std::atomic destroyed{0}; + + int value = kInitialValue; +}; +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); + +// Notifies `entered`, then waits for `gate`, for whichever is set. Lets a test +// keep a CowStorage insert, and so its mutex, held. +struct GatedState { + GatedState() { + if (absl::Notification* notification = entered.load()) { + notification->Notify(); + } + if (absl::Notification* notification = gate.load()) { + notification->WaitForNotification(); + } + } + + static void Reset() { + entered.store(nullptr); + gate.store(nullptr); + } + + static inline std::atomic entered{nullptr}; + static inline std::atomic gate{nullptr}; +}; + +bool IsCacheLineAligned(const void* ptr) { + return reinterpret_cast(ptr) % ABSL_CACHELINE_SIZE == 0; +} + +template +size_t CountDistinct(const std::vector& ptrs) { + return absl::flat_hash_set(ptrs.begin(), ptrs.end()).size(); +} + +class CowStorageTest : public ::testing::Test { + protected: + void SetUp() override { + Tracked::Reset(); + GatedState::Reset(); + } + void TearDown() override { GatedState::Reset(); } +}; + +TEST_F(CowStorageTest, FindReturnsNullForOrdinalsNotCreated) { + CowStorage storage; + EXPECT_EQ(storage.Find(-1), nullptr); + EXPECT_EQ(storage.GetOrCreate(-1, &DeviceSlot::Create), nullptr); + EXPECT_EQ(storage.Find(0), nullptr); + EXPECT_NE(storage.GetOrCreate(5, &DeviceSlot::Create), nullptr); + EXPECT_EQ(storage.Find(-1), nullptr); + EXPECT_EQ(storage.GetOrCreate(-1, &DeviceSlot::Create), nullptr); + EXPECT_EQ(storage.Find(4), nullptr); + EXPECT_EQ(storage.Find(6), nullptr); + EXPECT_EQ(Tracked::constructed.load(), 1); +} + +TEST_F(CowStorageTest, SparseOrdinals) { + CowStorage storage; + DeviceSlot* s0 = storage.GetOrCreate(0, &DeviceSlot::Create); + DeviceSlot* s7 = storage.GetOrCreate(7, &DeviceSlot::Create); + DeviceSlot* s31 = storage.GetOrCreate(31, &DeviceSlot::Create); + + EXPECT_NE(s0, nullptr); + EXPECT_NE(s7, nullptr); + EXPECT_NE(s31, nullptr); + EXPECT_EQ(storage.Find(0), s0); + EXPECT_EQ(storage.Find(7), s7); + EXPECT_EQ(storage.Find(31), s31); + for (int gap : {1, 6, 8, 30, 32}) { + EXPECT_EQ(storage.Find(gap), nullptr) << "ordinal " << gap; + } + EXPECT_EQ(Tracked::constructed.load(), 3); +} + +TEST_F(CowStorageTest, GetOrCreateBuildsOnce) { + CowStorage storage; + DeviceSlot* slot = storage.GetOrCreate(3, &DeviceSlot::Create); + EXPECT_NE(slot, nullptr); + EXPECT_EQ(storage.GetOrCreate(3, &DeviceSlot::Create), slot); + EXPECT_EQ(storage.Find(3), slot); + EXPECT_EQ(Tracked::constructed.load(), 1); +} + +TEST_F(CowStorageTest, StatesKeepAddressAcrossInserts) { + // More ordinals than a snapshot stores inline. + constexpr int kNumOrdinals = 16; + CowStorage storage; + std::vector states; + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + states.push_back(DeviceSlot::Unwrap( + storage.GetOrCreate(ordinal, &DeviceSlot::Create))); + } + EXPECT_THAT(states, Not(Contains(nullptr))); + EXPECT_EQ(CountDistinct(states), kNumOrdinals); + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + DeviceSlot* slot = storage.Find(ordinal); + EXPECT_EQ(DeviceSlot::Unwrap(slot), states[ordinal]) + << "ordinal " << ordinal; + EXPECT_TRUE(IsCacheLineAligned(slot)) << "ordinal " << ordinal; + } + EXPECT_EQ(Tracked::constructed.load(), kNumOrdinals); +} + +TEST_F(CowStorageTest, DestroysEveryState) { + { + CowStorage storage; + for (int ordinal = 0; ordinal < 3; ++ordinal) { + storage.GetOrCreate(ordinal, &DeviceSlot::Create); + } + EXPECT_EQ(Tracked::destroyed.load(), 0); + } + EXPECT_EQ(Tracked::constructed.load(), 3); + EXPECT_EQ(Tracked::destroyed.load(), 3); +} + +TEST_F(CowStorageTest, ConcurrentGetOrCreateBuildsOneState) { + constexpr int kNumThreads = 8; + CowStorage storage; + std::vector states(kNumThreads, nullptr); + absl::Notification start; + { + tsl::thread::ThreadPool pool(tsl::Env::Default(), "cow_storage_test", + kNumThreads); + for (int i = 0; i < kNumThreads; ++i) { + pool.Schedule([&, i] { + start.WaitForNotification(); + states[i] = DeviceSlot::Unwrap( + storage.GetOrCreate(0, &DeviceSlot::Create)); + }); + } + start.Notify(); + } // Joins the pool. + EXPECT_EQ(Tracked::constructed.load(), 1); + EXPECT_NE(states[0], nullptr); + EXPECT_THAT(states, Each(states[0])); +} + +TEST_F(CowStorageTest, ConcurrentReadersSeeStableFullyBuiltStates) { + constexpr int kNumOrdinals = 16; + constexpr int kNumReaders = 4; + CowStorage storage; + std::vector created(kNumOrdinals, nullptr); + std::vector> seen( + kNumReaders, std::vector(kNumOrdinals, nullptr)); + absl::Notification start; + std::atomic writers_done{false}; + std::atomic errors{0}; + { + tsl::thread::ThreadPool readers(tsl::Env::Default(), "readers", + kNumReaders); + for (int reader = 0; reader < kNumReaders; ++reader) { + readers.Schedule([&, reader] { + start.WaitForNotification(); + std::vector& last_seen = seen[reader]; + bool done = false; + while (!done) { + // Loaded before the scan, so the last scan starts after all writers + // have finished. + done = writers_done.load(); + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + Tracked* state = DeviceSlot::Unwrap(storage.Find(ordinal)); + if (state == nullptr) { + continue; + } + // A published state is fully built and never replaced. + if (state->value != Tracked::kInitialValue || + (last_seen[ordinal] != nullptr && + last_seen[ordinal] != state)) { + errors.fetch_add(1); + } + last_seen[ordinal] = state; + } + } + }); + } + { + tsl::thread::ThreadPool writers(tsl::Env::Default(), "writers", + kNumOrdinals); + for (int ordinal = 0; ordinal < kNumOrdinals; ++ordinal) { + writers.Schedule([&, ordinal] { + start.WaitForNotification(); + created[ordinal] = DeviceSlot::Unwrap( + storage.GetOrCreate(ordinal, &DeviceSlot::Create)); + }); + } + start.Notify(); + } // Joins the writers. + writers_done.store(true); + } // Joins the readers. + EXPECT_EQ(errors.load(), 0); + EXPECT_EQ(Tracked::constructed.load(), kNumOrdinals); + EXPECT_THAT(created, Not(Contains(nullptr))); + for (int reader = 0; reader < kNumReaders; ++reader) { + EXPECT_EQ(seen[reader], created) << "reader " << reader; + } +} + +// An insert holds the mutex while it builds the state. `Find` must not wait +// for it. +TEST_F(CowStorageTest, FindDoesNotWaitForInsert) { + CowStorage storage; + DeviceSlot* existing = + storage.GetOrCreate(0, &DeviceSlot::Create); + ASSERT_NE(existing, nullptr); + + absl::Notification entered; + absl::Notification release; + absl::Notification read_done; + GatedState::entered.store(&entered); + GatedState::gate.store(&release); + + bool insert_in_progress = false; + bool read_finished = false; + DeviceSlot* found = nullptr; + DeviceSlot* uncommitted = nullptr; + { + tsl::thread::ThreadPool pool(tsl::Env::Default(), "cow_storage_test", 2); + pool.Schedule( + [&] { storage.GetOrCreate(1, &DeviceSlot::Create); }); + insert_in_progress = + entered.WaitForNotificationWithTimeout(absl::Seconds(10)); + if (insert_in_progress) { + pool.Schedule([&] { + found = storage.Find(0); + uncommitted = storage.Find(1); + read_done.Notify(); + }); + read_finished = + read_done.WaitForNotificationWithTimeout(absl::Seconds(10)); + } + // Always unblock the insert, so the pool can join. + release.Notify(); + } // Joins the pool. + GatedState::Reset(); + + EXPECT_TRUE(insert_in_progress); + EXPECT_TRUE(read_finished); + EXPECT_EQ(found, existing); + EXPECT_EQ(uncommitted, nullptr); +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/device_slot.h b/third_party/xla/xla/backends/gpu/runtime/device_slot.h new file mode 100644 index 00000000000000..02640185ab6f97 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/device_slot.h @@ -0,0 +1,59 @@ +/* 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_BACKENDS_GPU_RUNTIME_DEVICE_SLOT_H_ +#define XLA_BACKENDS_GPU_RUNTIME_DEVICE_SLOT_H_ + +#include + +#include "absl/base/optimization.h" +#include "absl/functional/function_ref.h" + +namespace xla::gpu { + +// Type-erased, cache-line-aligned base for a per-device state slot. +// Aligned so that states of different devices never share a cache line. +struct alignas(ABSL_CACHELINE_SIZE) DeviceSlot { + virtual ~DeviceSlot() = default; + + template + static std::unique_ptr Create(); + + template + static T* Unwrap(DeviceSlot* slot); +}; + +template +struct TypedDeviceSlot final : DeviceSlot { + T value; +}; + +template +std::unique_ptr DeviceSlot::Create() { + return std::make_unique>(); +} + +template +T* DeviceSlot::Unwrap(DeviceSlot* slot) { + return slot != nullptr ? &(static_cast*>(slot)->value) + : nullptr; +} + +// No argument function that creates a DeviceSlot. +using DeviceSlotFactoryRef = absl::FunctionRef()>; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_RUNTIME_DEVICE_SLOT_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/vector_storage.cc b/third_party/xla/xla/backends/gpu/runtime/vector_storage.cc new file mode 100644 index 00000000000000..d0f3b2112db59b --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/vector_storage.cc @@ -0,0 +1,29 @@ +/* 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/backends/gpu/runtime/vector_storage.h" + +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +DeviceSlot* VectorStorage::Find(int device_ordinal) const { + if (device_ordinal < 0 || device_ordinal >= size()) { + return nullptr; + } + return slots_[device_ordinal].get(); +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/vector_storage.h b/third_party/xla/xla/backends/gpu/runtime/vector_storage.h new file mode 100644 index 00000000000000..3be2705cab164a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/vector_storage.h @@ -0,0 +1,49 @@ +/* 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_BACKENDS_GPU_RUNTIME_VECTOR_STORAGE_H_ +#define XLA_BACKENDS_GPU_RUNTIME_VECTOR_STORAGE_H_ + +#include +#include +#include + +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { + +// Pre-sized storage for holding per-device state slots in [0, size), which can +// be accessed across threads safely as long as each thread only accesses a +// single ordinal/index. `Find` is a bounds check plus an array index. +class VectorStorage { + public: + explicit VectorStorage(std::vector> slots) + : slots_(std::move(slots)) {} + + VectorStorage(const VectorStorage&) = delete; + VectorStorage& operator=(const VectorStorage&) = delete; + + int size() const { return static_cast(slots_.size()); } + + // Returns nullptr if `device_ordinal` is outside [0, size). + DeviceSlot* Find(int device_ordinal) const; + + private: + std::vector> slots_; +}; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_RUNTIME_VECTOR_STORAGE_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc b/third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc new file mode 100644 index 00000000000000..2ee5c802fae2ef --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/vector_storage_test.cc @@ -0,0 +1,138 @@ +/* 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/backends/gpu/runtime/vector_storage.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include "absl/base/optimization.h" +#include "absl/container/flat_hash_set.h" +#include "xla/backends/gpu/runtime/device_slot.h" + +namespace xla::gpu { +namespace { + +using ::testing::Contains; +using ::testing::Not; + +// Counts constructions and destructions. Neither copyable nor movable, so a +// storage that copies or moves its states fails to compile. +struct Tracked { + static constexpr int kInitialValue = 42; + + Tracked() { constructed.fetch_add(1); } + ~Tracked() { destroyed.fetch_add(1); } + Tracked(const Tracked&) = delete; + Tracked& operator=(const Tracked&) = delete; + + static void Reset() { + constructed.store(0); + destroyed.store(0); + } + + static inline std::atomic constructed{0}; + static inline std::atomic destroyed{0}; + + int value = kInitialValue; +}; +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(!std::is_move_constructible_v); +static_assert(!std::is_move_assignable_v); + +template +std::vector> MakeSlots(int size) { + std::vector> slots; + slots.reserve(size); + for (int i = 0; i < size; ++i) { + slots.push_back(DeviceSlot::Create()); + } + return slots; +} + +bool IsCacheLineAligned(const void* ptr) { + return reinterpret_cast(ptr) % ABSL_CACHELINE_SIZE == 0; +} + +template +size_t CountDistinct(const std::vector& ptrs) { + return absl::flat_hash_set(ptrs.begin(), ptrs.end()).size(); +} + +class VectorStorageTest : public ::testing::Test { + protected: + void SetUp() override { Tracked::Reset(); } +}; + +TEST_F(VectorStorageTest, StatesExistAfterConstruction) { + { + VectorStorage storage(MakeSlots(4)); + EXPECT_EQ(storage.size(), 4); + EXPECT_EQ(Tracked::constructed.load(), 4); + for (int ordinal = 0; ordinal < 4; ++ordinal) { + Tracked* state = DeviceSlot::Unwrap(storage.Find(ordinal)); + ASSERT_NE(state, nullptr) << "ordinal " << ordinal; + EXPECT_EQ(state->value, Tracked::kInitialValue) << "ordinal " << ordinal; + } + } + EXPECT_EQ(Tracked::destroyed.load(), 4); +} + +TEST_F(VectorStorageTest, FindOutsideRangeReturnsNull) { + VectorStorage storage(MakeSlots(4)); + EXPECT_EQ(storage.Find(-1), nullptr); + EXPECT_EQ(storage.Find(4), nullptr); +} + +TEST_F(VectorStorageTest, EmptyStorageFindsNothing) { + VectorStorage empty({}); + EXPECT_EQ(empty.size(), 0); + EXPECT_EQ(empty.Find(0), nullptr); + EXPECT_EQ(Tracked::constructed.load(), 0); +} + +TEST_F(VectorStorageTest, StatesAreDistinctAndStable) { + VectorStorage storage(MakeSlots(4)); + std::vector states; + for (int ordinal = 0; ordinal < 4; ++ordinal) { + states.push_back(DeviceSlot::Unwrap(storage.Find(ordinal))); + } + EXPECT_THAT(states, Not(Contains(nullptr))); + EXPECT_EQ(CountDistinct(states), 4); + for (int ordinal = 0; ordinal < 4; ++ordinal) { + EXPECT_EQ(DeviceSlot::Unwrap(storage.Find(ordinal)), + states[ordinal]) + << "ordinal " << ordinal; + } +} + +TEST_F(VectorStorageTest, StatesAreCacheLineAligned) { + VectorStorage storage(MakeSlots(4)); + for (int ordinal = 0; ordinal < 4; ++ordinal) { + DeviceSlot* slot = storage.Find(ordinal); + EXPECT_NE(slot, nullptr) << "ordinal " << ordinal; + EXPECT_TRUE(IsCacheLineAligned(slot)) << "ordinal " << ordinal; + } +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 73e0ea8e1b3810..c7fd9b12c2af57 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -296,6 +296,7 @@ xla_test( deps = [ ":hlo_pjrt_gpu_test_base", "//xla:error_spec", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:filecheck", "//xla/stream_executor:device_description", @@ -1171,6 +1172,7 @@ cc_library( "//xla/stream_executor/cuda:cuda_compute_capability", "//xla/tsl/platform:env", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", @@ -1712,6 +1714,8 @@ xla_test( "//xla/tests:literal_test_util", "//xla/tsl/platform:test", "//xla/tsl/testing:temporary_directory", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc index 0a118b107155db..60c7b25fb614ec 100644 --- a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc @@ -13,12 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include #include #include +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "xla/backends/gpu/tests/collective_ops_e2e_test_base.h" #include "xla/hlo/ir/hlo_computation.h" @@ -251,9 +254,10 @@ TEST_F(AllGatherTest, TwoDimensional2Gpu) { } } -// Runs 10 consequent all-gathers with FUSION command buffers enabled and -// verifies that all of them use XLA's one-shot all-gather kernel and are placed -// into command buffers. +// Runs 10 consequent all-gathers and verifies that by default all of them use +// XLA's one-shot all-gather kernel and are placed into command buffers, and +// that excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer +// disables command buffer capture for them. TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { constexpr int32_t kNumReplicas = 2; constexpr int kNumAllGathers = 10; @@ -281,6 +285,11 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { ASSERT_OK_AND_ASSIGN( tsl::testing::TemporaryDirectory dump_dir, tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + const std::string default_dump_dir = + absl::StrCat(dump_dir.path(), "/default"); + const std::string no_graph_dump_dir = + absl::StrCat(dump_dir.path(), "/no_graphs"); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); @@ -288,15 +297,13 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { module->mutable_config().mutable_debug_options(); // The fixture enables COLLECTIVE_KERNEL_ALL_GATHER, so the 10 all-gathers // lower to XLA's Triton one-shot all-gather kernel (CollectiveKernelThunk) - // instead of NCCL. Collective kernels are recorded into command buffers - // with the FUSION command buffer type, so COLLECTIVES isn't needed. - debug_options.clear_xla_gpu_enable_command_buffer(); - debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::COLLECTIVES); + // instead of NCCL. COLLECTIVES_KERNEL is enabled in + // --xla_gpu_enable_command_buffer by default. debug_options.set_xla_gpu_graph_min_graph_size(1); debug_options.set_xla_gpu_all_gather_combine_threshold_bytes(0); // Command buffers are formed at the thunk level, so dump the thunk sequence // to check which collectives were placed into them. - debug_options.set_xla_dump_to(dump_dir.path()); + debug_options.set_xla_dump_to(default_dump_dir); Literal input_r0 = LiteralUtil::CreateR1(std::vector(16, 1.0f)); Literal input_r1 = LiteralUtil::CreateR1(std::vector(16, 2.0f)); @@ -308,14 +315,15 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { // kernel, not NCCL). VerifyOneShotAllGather(result.optimized_module); - // All one-shot all-gathers must be recorded into command buffers, and none - // may fall back to NCCL. - ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts one_shot, - CountThunksInDump(dump_dir.path(), "kCollectiveKernel")); + // By default, all one-shot all-gathers must be recorded into command buffers, + // and none may fall back to NCCL. + ASSERT_OK_AND_ASSIGN( + CommandBufferThunkCounts one_shot, + CountThunksInDump(default_dump_dir, "kCollectiveKernel")); EXPECT_EQ(one_shot.in_command_buffer, kNumAllGathers); EXPECT_EQ(one_shot.outside_command_buffer, 0); ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl, - CountThunksInDump(dump_dir.path(), "kAllGather")); + CountThunksInDump(default_dump_dir, "kAllGather")); EXPECT_EQ(nccl.in_command_buffer + nccl.outside_command_buffer, 0); // Each all-gather concatenates the two ranks' buffers, so the result is 512 @@ -332,6 +340,173 @@ TEST_F(AllGatherTest, TenConsequentAllGathersWithCudaGraphs) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) << "Mismatch at replica " << i; } + + // Excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer must + // disable command buffer capture for the collective kernels. + ASSERT_OK_AND_ASSIGN(auto no_graph_module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + DebugOptions& no_graph_debug_options = + no_graph_module->mutable_config().mutable_debug_options(); + auto* enabled_commands = + no_graph_debug_options.mutable_xla_gpu_enable_command_buffer(); + enabled_commands->erase( + std::remove(enabled_commands->begin(), enabled_commands->end(), + DebugOptions::COLLECTIVES_KERNEL), + enabled_commands->end()); + no_graph_debug_options.set_xla_gpu_graph_min_graph_size(1); + no_graph_debug_options.set_xla_gpu_all_gather_combine_threshold_bytes(0); + no_graph_debug_options.set_xla_dump_to(no_graph_dump_dir); + + ASSERT_OK_AND_ASSIGN(ExecutionResult no_graph_result, + ExecuteReplicated(std::move(no_graph_module), args)); + VerifyOneShotAllGather(no_graph_result.optimized_module); + + ASSERT_OK_AND_ASSIGN( + CommandBufferThunkCounts one_shot_no_graph, + CountThunksInDump(no_graph_dump_dir, "kCollectiveKernel")); + EXPECT_EQ(one_shot_no_graph.in_command_buffer, 0); + EXPECT_EQ(one_shot_no_graph.outside_command_buffer, kNumAllGathers); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl_no_graph, + CountThunksInDump(no_graph_dump_dir, "kAllGather")); + EXPECT_EQ( + nccl_no_graph.in_command_buffer + nccl_no_graph.outside_command_buffer, + 0); + ASSERT_EQ(no_graph_result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, no_graph_result.results[i])) + << "Mismatch at replica " << i << " without graphs"; + } +} + +// Runs 20 one-shot all-gathers inside a 20-iteration while loop with WHILE +// command buffers enabled, and verifies that the while thunk and all 20 +// collective kernels are placed into command buffers and produce the expected +// result across repeated executions (exercising the persistent device-side +// invocation counter). +TEST_F(AllGatherTest, TwentyAllGathersInWhileLoopWithCudaGraphs) { + constexpr int32_t kNumReplicas = 2; + constexpr int kNumAllGathers = 20; + constexpr int kNumIterations = 20; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + // Each step k in the loop body gathers in_k = 0.5 * prev + p0 (f32[128]) + // across the 2 replicas into ag_k (f32[256]) and sums the two 128-element + // slices: (0.5 * v + 1.0) + (0.5 * v + 2.0) = v + 3.0. Every all-gather + // therefore depends on both replicas' slices from the previous all-gather and + // adds 3.0 without growing the tensor shape or overflowing f32 across + // 20 * 20 = 400 collective launches. + std::string body_ops; + std::string prev = "acc"; + for (int k = 0; k < kNumAllGathers; ++k) { + absl::StrAppendFormat( + &body_ops, + " scaled%d = f32[128] multiply(%s, half)\n" + " in%d = f32[128] add(scaled%d, p0)\n" + " ag%d = f32[256] all-gather(in%d), dimensions={0}, " + "replica_groups={{0,1}}\n" + " s0_%d = f32[128] slice(ag%d), slice={[0:128]}\n" + " s1_%d = f32[128] slice(ag%d), slice={[128:256]}\n" + " sum%d = f32[128] add(s0_%d, s1_%d)\n", + k, prev, k, k, k, k, k, k, k, k, k, k, k); + prev = absl::StrCat("sum", k); + } + + const std::string module_str = absl::StrFormat( + R"( + HloModule test + + cond { + state = (s32[], f32[128], f32[128]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + limit = s32[] constant(%d) + ROOT cmp = pred[] compare(i, limit), direction=LT + } + + body { + state = (s32[], f32[128], f32[128]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + acc = f32[128] get-tuple-element(state), index=1 + p0 = f32[128] get-tuple-element(state), index=2 + one = s32[] constant(1) + next_i = s32[] add(i, one) + half_scalar = f32[] constant(0.5) + half = f32[128] broadcast(half_scalar), dimensions={} +%s ROOT next_state = (s32[], f32[128], f32[128]) tuple(next_i, %s, p0) + } + + ENTRY test_computation { + p0 = f32[128] parameter(0) + zero_i = s32[] constant(0) + zero_f = f32[] constant(0.0) + init_acc = f32[128] broadcast(zero_f), dimensions={} + init = (s32[], f32[128], f32[128]) tuple(zero_i, init_acc, p0) + loop = (s32[], f32[128], f32[128]) while(init), condition=cond, body=body + ROOT out = f32[128] get-tuple-element(loop), index=1 + } + )", + kNumIterations, body_ops, prev); + + ASSERT_OK_AND_ASSIGN( + tsl::testing::TemporaryDirectory dump_dir, + tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(module_str, kNumReplicas)); + + DebugOptions& debug_options = + module->mutable_config().mutable_debug_options(); + debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::WHILE); + debug_options.set_xla_gpu_graph_min_graph_size(1); + debug_options.set_xla_gpu_all_gather_combine_threshold_bytes(0); + debug_options.set_xla_gpu_enable_while_loop_unrolling( + DebugOptions::WHILE_LOOP_UNROLLING_NO_UNROLL); + debug_options.set_xla_dump_to(dump_dir.path()); + + Literal input_r0 = + LiteralUtil::CreateR1(std::vector(128, 1.0f)); + Literal input_r1 = + LiteralUtil::CreateR1(std::vector(128, 2.0f)); + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + VerifyOneShotAllGather(result.optimized_module); + + // Both the while thunk and all 20 one-shot all-gathers in its body must be + // recorded into command buffers, and none may fall back to NCCL. + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts while_counts, + CountThunksInDump(dump_dir.path(), "kWhile")); + EXPECT_EQ(while_counts.in_command_buffer, 1); + EXPECT_EQ(while_counts.outside_command_buffer, 0); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts one_shot, + CountThunksInDump(dump_dir.path(), "kCollectiveKernel")); + EXPECT_EQ(one_shot.in_command_buffer, kNumAllGathers); + EXPECT_EQ(one_shot.outside_command_buffer, 0); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl, + CountThunksInDump(dump_dir.path(), "kAllGather")); + EXPECT_EQ(nccl.in_command_buffer + nccl.outside_command_buffer, 0); + + // 20 iterations * 20 all-gathers * 3.0 per all-gather = 1200.0. + const float expected_val = 3.0f * kNumAllGathers * kNumIterations; + Literal expected = + LiteralUtil::CreateR1(std::vector(128, expected_val)); + ASSERT_EQ(result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i << " on first execution"; + } + + // Re-run the same executable to verify that the device-side counters persist + // across executions and the replayed CUDA graph still produces the expected + // result. + ASSERT_OK_AND_ASSIGN(std::vector second_results, + ExecuteReplicated(result.executable.get(), args)); + ASSERT_EQ(second_results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, second_results[i])) + << "Mismatch at replica " << i << " on second execution"; + } } } // namespace diff --git a/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc index fa2089ef918828..97780efd450c51 100644 --- a/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/all_reduce_e2e_test.cc @@ -1101,9 +1101,10 @@ TEST_F(AllReduceTestNoParams, TestGlobalDeviceIdMappingWithAsyncAllReduce) { } } -// Runs 10 consequent all-reduces with FUSION command buffers enabled and -// verifies that all of them use XLA's one-shot all-reduce kernel and are placed -// into command buffers. +// Runs 10 consequent all-reduces and verifies that by default all of them use +// XLA's one-shot all-reduce kernel and are placed into command buffers, and +// that excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer +// disables command buffer capture for them. TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { constexpr int64_t kNumReplicas = 2; constexpr int kNumAllReduces = 10; @@ -1158,6 +1159,11 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { ASSERT_OK_AND_ASSIGN( tsl::testing::TemporaryDirectory dump_dir, tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + const std::string default_dump_dir = + absl::StrCat(dump_dir.path(), "/default"); + const std::string no_graph_dump_dir = + absl::StrCat(dump_dir.path(), "/no_graphs"); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); @@ -1165,15 +1171,13 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { module->mutable_config().mutable_debug_options(); // The fixture enables COLLECTIVE_KERNEL_ALL_REDUCE, so the 10 all-reduces // lower to XLA's Triton one-shot all-reduce kernel (CollectiveKernelThunk) - // instead of NCCL. Collective kernels are recorded into command buffers - // with the FUSION command buffer type, so COLLECTIVES isn't needed. - debug_options.clear_xla_gpu_enable_command_buffer(); - debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::COLLECTIVES); + // instead of NCCL. COLLECTIVES_KERNEL is enabled in + // --xla_gpu_enable_command_buffer by default. debug_options.set_xla_gpu_graph_min_graph_size(1); debug_options.set_xla_gpu_all_reduce_combine_threshold_bytes(0); // Command buffers are formed at the thunk level, so dump the thunk sequence // to check which collectives were placed into them. - debug_options.set_xla_dump_to(dump_dir.path()); + debug_options.set_xla_dump_to(default_dump_dir); Literal input_r0 = LiteralUtil::CreateR1(std::vector(1024, 1.0f)); @@ -1187,14 +1191,15 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { // kernel, not NCCL). VerifyOneShotAllReduce(result.optimized_module); - // All one-shot all-reduces must be recorded into command buffers, and none - // may fall back to NCCL. - ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts one_shot, - CountThunksInDump(dump_dir.path(), "kCollectiveKernel")); + // By default, all one-shot all-reduces must be recorded into command buffers, + // and none may fall back to NCCL. + ASSERT_OK_AND_ASSIGN( + CommandBufferThunkCounts one_shot, + CountThunksInDump(default_dump_dir, "kCollectiveKernel")); EXPECT_EQ(one_shot.in_command_buffer, kNumAllReduces); EXPECT_EQ(one_shot.outside_command_buffer, 0); ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl, - CountThunksInDump(dump_dir.path(), "kAllReduce")); + CountThunksInDump(default_dump_dir, "kAllReduce")); EXPECT_EQ(nccl.in_command_buffer + nccl.outside_command_buffer, 0); // The first all-reduce sums 1 + 2 = 3 and each following one computes @@ -1207,6 +1212,175 @@ TEST_F(AllReduceCollectiveKernelTest, TenConsequentAllReducesWithCudaGraphs) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) << "Mismatch at replica " << i; } + + // Excluding COLLECTIVES_KERNEL from --xla_gpu_enable_command_buffer must + // disable command buffer capture for the collective kernels. + ASSERT_OK_AND_ASSIGN(auto no_graph_module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + DebugOptions& no_graph_debug_options = + no_graph_module->mutable_config().mutable_debug_options(); + auto* enabled_commands = + no_graph_debug_options.mutable_xla_gpu_enable_command_buffer(); + enabled_commands->erase( + std::remove(enabled_commands->begin(), enabled_commands->end(), + DebugOptions::COLLECTIVES_KERNEL), + enabled_commands->end()); + no_graph_debug_options.set_xla_gpu_graph_min_graph_size(1); + no_graph_debug_options.set_xla_gpu_all_reduce_combine_threshold_bytes(0); + no_graph_debug_options.set_xla_dump_to(no_graph_dump_dir); + + ASSERT_OK_AND_ASSIGN(ExecutionResult no_graph_result, + ExecuteReplicated(std::move(no_graph_module), args)); + VerifyOneShotAllReduce(no_graph_result.optimized_module); + + ASSERT_OK_AND_ASSIGN( + CommandBufferThunkCounts one_shot_no_graph, + CountThunksInDump(no_graph_dump_dir, "kCollectiveKernel")); + EXPECT_EQ(one_shot_no_graph.in_command_buffer, 0); + EXPECT_EQ(one_shot_no_graph.outside_command_buffer, kNumAllReduces); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl_no_graph, + CountThunksInDump(no_graph_dump_dir, "kAllReduce")); + EXPECT_EQ( + nccl_no_graph.in_command_buffer + nccl_no_graph.outside_command_buffer, + 0); + ASSERT_EQ(no_graph_result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, no_graph_result.results[i])) + << "Mismatch at replica " << i << " without graphs"; + } +} + +// Runs 20 one-shot all-reduces inside a 20-iteration while loop with WHILE +// command buffers enabled, and verifies that the while thunk and all 20 +// collective kernels are placed into command buffers and produce the expected +// result across repeated executions (exercising the persistent device-side +// invocation counter). +TEST_F(AllReduceCollectiveKernelTest, + TwentyAllReducesInWhileLoopWithCudaGraphs) { + constexpr int64_t kNumReplicas = 2; + constexpr int kNumAllReduces = 20; + constexpr int kNumIterations = 20; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + // Each step k in the loop body computes ar_k = all-reduce(0.5 * prev + p0), + // where p0 is 1.0 on replica 0 and 2.0 on replica 1. When prev == v on both + // replicas, (0.5 * v + 1.0) + (0.5 * v + 2.0) = v + 3.0, so every all-reduce + // adds 3.0 without overflowing f32 across 20 * 20 = 400 collective launches. + std::string body_ops; + std::string prev = "acc"; + for (int k = 0; k < kNumAllReduces; ++k) { + absl::StrAppendFormat( + &body_ops, + " scaled%d = f32[1024] multiply(%s, half)\n" + " in%d = f32[1024] add(scaled%d, p0)\n" + " ar%d = f32[1024] all-reduce(in%d), to_apply=add, " + "replica_groups={{0,1}}\n", + k, prev, k, k, k, k); + prev = absl::StrCat("ar", k); + } + + const std::string module_str = absl::StrFormat( + R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT sum = f32[] add(x, y) + } + + cond { + state = (s32[], f32[1024], f32[1024]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + limit = s32[] constant(%d) + ROOT cmp = pred[] compare(i, limit), direction=LT + } + + body { + state = (s32[], f32[1024], f32[1024]) parameter(0) + i = s32[] get-tuple-element(state), index=0 + acc = f32[1024] get-tuple-element(state), index=1 + p0 = f32[1024] get-tuple-element(state), index=2 + one = s32[] constant(1) + next_i = s32[] add(i, one) + half_scalar = f32[] constant(0.5) + half = f32[1024] broadcast(half_scalar), dimensions={} +%s ROOT next_state = (s32[], f32[1024], f32[1024]) tuple(next_i, %s, p0) + } + + ENTRY test_computation { + p0 = f32[1024] parameter(0) + zero_i = s32[] constant(0) + zero_f = f32[] constant(0.0) + init_acc = f32[1024] broadcast(zero_f), dimensions={} + init = (s32[], f32[1024], f32[1024]) tuple(zero_i, init_acc, p0) + loop = (s32[], f32[1024], f32[1024]) while(init), condition=cond, body=body + ROOT out = f32[1024] get-tuple-element(loop), index=1 + } + )", + kNumIterations, body_ops, prev); + + ASSERT_OK_AND_ASSIGN( + tsl::testing::TemporaryDirectory dump_dir, + tsl::testing::TemporaryDirectory::CreateForCurrentTestcase()); + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(module_str, kNumReplicas)); + + DebugOptions& debug_options = + module->mutable_config().mutable_debug_options(); + debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::WHILE); + debug_options.set_xla_gpu_graph_min_graph_size(1); + debug_options.set_xla_gpu_all_reduce_combine_threshold_bytes(0); + debug_options.set_xla_gpu_enable_while_loop_unrolling( + DebugOptions::WHILE_LOOP_UNROLLING_NO_UNROLL); + debug_options.set_xla_dump_to(dump_dir.path()); + + Literal input_r0 = + LiteralUtil::CreateR1(std::vector(1024, 1.0f)); + Literal input_r1 = + LiteralUtil::CreateR1(std::vector(1024, 2.0f)); + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + VerifyOneShotAllReduce(result.optimized_module); + + // Both the while thunk and all 20 one-shot all-reduces in its body must be + // recorded into command buffers, and none may fall back to NCCL. + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts while_counts, + CountThunksInDump(dump_dir.path(), "kWhile")); + EXPECT_EQ(while_counts.in_command_buffer, 1); + EXPECT_EQ(while_counts.outside_command_buffer, 0); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts one_shot, + CountThunksInDump(dump_dir.path(), "kCollectiveKernel")); + EXPECT_EQ(one_shot.in_command_buffer, kNumAllReduces); + EXPECT_EQ(one_shot.outside_command_buffer, 0); + ASSERT_OK_AND_ASSIGN(CommandBufferThunkCounts nccl, + CountThunksInDump(dump_dir.path(), "kAllReduce")); + EXPECT_EQ(nccl.in_command_buffer + nccl.outside_command_buffer, 0); + + // 20 iterations * 20 all-reduces * 3.0 per all-reduce = 1200.0. + const float expected_val = 3.0f * kNumAllReduces * kNumIterations; + Literal expected = + LiteralUtil::CreateR1(std::vector(1024, expected_val)); + ASSERT_EQ(result.results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i << " on first execution"; + } + + // Re-run the same executable to verify that the device-side counters persist + // across executions and the replayed CUDA graph still produces the expected + // result. + ASSERT_OK_AND_ASSIGN(std::vector second_results, + ExecuteReplicated(result.executable.get(), args)); + ASSERT_EQ(second_results.size(), kNumReplicas); + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, second_results[i])) + << "Mismatch at replica " << i << " on second execution"; + } } } // namespace diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc index ab11d842402e53..0a61389bad351c 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test_base.cc @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" @@ -204,6 +205,12 @@ absl::StatusOr CountThunksInDump( ABSL_RETURN_IF_ERROR(tsl::Env::Default()->GetMatchingPaths( tsl::io::JoinPath(dump_dir, "*thunk_sequence_after_thunk_passes*.txt"), &dump_files)); + if (dump_files.empty()) { + // When thunk passes make no changes (e.g. no command buffers are formed), + // only the initial thunk_sequence.txt dump is written. + ABSL_RETURN_IF_ERROR(tsl::Env::Default()->GetMatchingPaths( + tsl::io::JoinPath(dump_dir, "*thunk_sequence.txt"), &dump_files)); + } if (dump_files.size() != 1) { return absl::FailedPreconditionError( absl::StrCat("Expected exactly one thunk sequence dump in ", dump_dir, @@ -213,19 +220,22 @@ absl::StatusOr CountThunksInDump( ABSL_RETURN_IF_ERROR( tsl::ReadFileToString(tsl::Env::Default(), dump_files[0], &dump)); - // Each line of the dump describes one thunk as ": ...". - // Thunks nested in a kCommandBuffer thunk are indented under it. + // Each thunk line in the dump has the form ": ...". + // Thunks nested inside a top-level kCommandBuffer thunk are indented under + // it; non-thunk lines (such as WhileThunk's "condition:" and "body:" headers) + // do not start with ": " and are ignored. CommandBufferThunkCounts counts; bool in_command_buffer = false; for (absl::string_view line : absl::StrSplit(dump, '\n')) { - absl::string_view thunk = absl::StripAsciiWhitespace(line); - if (thunk.empty()) { + absl::string_view thunk = absl::StripLeadingAsciiWhitespace(line); + size_t pos = thunk.find(": "); + if (pos == absl::string_view::npos || pos == 0 || + !absl::c_all_of(thunk.substr(0, pos), + [](char c) { return absl::ascii_isdigit(c); })) { continue; } - const bool is_top_level = thunk.data() == line.data(); - if (size_t pos = thunk.find(": "); pos != absl::string_view::npos) { - thunk.remove_prefix(pos + 2); - } + const bool is_top_level = thunk.size() == line.size(); + thunk.remove_prefix(pos + 2); if (is_top_level) { in_command_buffer = absl::StartsWith(thunk, "kCommandBuffer"); } diff --git a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc index 7ed04346490c4f..941aef2f58280f 100644 --- a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "xla/hlo/testlib/filecheck.h" #include "xla/stream_executor/device_description.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" +#include "xla/xla.pb.h" namespace xla::gpu { namespace { @@ -33,6 +34,13 @@ namespace { class SwapConvOperandsTest : public HloInterpreterReferenceMixin { public: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = HloInterpreterReferenceMixin< + HloPjRtGpuTestBase>::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_enable_conv_fusion(false); + return debug_options; + } + void MatchOptimizedHlo(absl::string_view hlo, absl::string_view expected_hlo) { ASSERT_OK_AND_ASSIGN(std::unique_ptr optimized_module, diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index c78569045a3131..4c4089e2c614d0 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -2319,8 +2319,8 @@ cc_library( "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/service:computation_layout", + "//xla/service:hlo_value", "//xla/service:layout_assignment", - "//xla/service:logical_buffer", "//xla/service:matmul_indexing_utils", "//xla/service:memory_annotations_hdr", "//xla/service/gpu:conv_utils", diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc index 1304cdeb81d8a7..027abcc82d77a6 100644 --- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc @@ -268,9 +268,10 @@ absl::StatusOr ProcessFusionInstruction( auto fusion_adaptor = HloFusionAdaptor::ForInstruction( Cast(fusion_instruction)); - ABSL_ASSIGN_OR_RETURN( - TiledRunTimeDataOrError tiled_runtime_data_or_error, - indexing_performance_model.TryFindBestTilingForFusion(*fusion_adaptor)); + ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or_error, + indexing_performance_model + .TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* fusion_decision = std::get_if(&tiled_runtime_data_or_error)) { diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc index a1f641becd85d1..ff0811a8a21a31 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc @@ -50,8 +50,8 @@ limitations under the License. #include "xla/service/gpu/matmul_utils.h" #include "xla/service/gpu/reduction_utils.h" #include "xla/service/gpu/stream_executor_util.h" +#include "xla/service/hlo_value.h" #include "xla/service/layout_assignment.h" -#include "xla/service/logical_buffer.h" #include "xla/service/matmul_indexing_utils.h" #include "xla/service/memory_annotations.h" #include "xla/shape.h" @@ -253,14 +253,13 @@ absl::Status GpuLayoutAssignment::AddBackendConstraintsToDnnConvCustomCall( // The custom call returns a tuple of (actual_result, scratch_buffer); // call_result_buf is the logical buffer for actual_result, the thing that // contains the result of the conv call. - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* call_result_buf, - points_to_analysis_->GetBufferDefinedAt(instr, /*index=*/{0})); + const HloValue& call_result_buf = + dataflow_analysis_->GetValueDefinedAt(instr, /*index=*/{0}); // Set layouts of the instructions' shapes. ABSL_RETURN_IF_ERROR(SetOperandLayout(lhs_shape, instr, 0)); ABSL_RETURN_IF_ERROR(SetOperandLayout(rhs_shape, instr, 1)); - ABSL_RETURN_IF_ERROR(SetBufferLayout(result_shape.layout(), *call_result_buf)); + ABSL_RETURN_IF_ERROR(SetBufferLayout(result_shape.layout(), call_result_buf)); // For fused convolutions, instr->operand(2), if exists, is the bias buffer. // There is no need to assign layout to it, as it has only one dimension. // instr->operand(3), if exists, is the side input buffer. @@ -648,30 +647,22 @@ absl::Status GpuLayoutAssignment::AddBackendConstraints( Shape shape = instruction->operand(i)->shape(); *shape.mutable_layout() = keys_layout; ABSL_RETURN_IF_ERROR(SetOperandLayout(shape, instruction, i)); - const LogicalBuffer* output_buffer; - if (instruction->shape().IsArray()) { - ABSL_ASSIGN_OR_RETURN( - output_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {})); - } else { - ABSL_ASSIGN_OR_RETURN( - output_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {i})); - } - ABSL_RETURN_IF_ERROR(SetBufferLayout(keys_layout, *output_buffer)); + const HloValue& output_buffer = + instruction->shape().IsArray() + ? dataflow_analysis_->GetValueDefinedAt(instruction, {}) + : dataflow_analysis_->GetValueDefinedAt(instruction, {i}); + ABSL_RETURN_IF_ERROR(SetBufferLayout(keys_layout, output_buffer)); } } else if (IsCustomCallToTopK(*instruction)) { // The output of the TopK custom call needs to have default layout. Layout default_layout = LayoutUtil::GetDefaultLayoutForRank( instruction->operand(0)->shape().dimensions().size()); - ABSL_ASSIGN_OR_RETURN( - auto values_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {0})); - ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, *values_buffer)); - ABSL_ASSIGN_OR_RETURN( - auto indices_buffer, - points_to_analysis_->GetBufferDefinedAt(instruction, {1})); - ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, *indices_buffer)); + const HloValue& values_buffer = + dataflow_analysis_->GetValueDefinedAt(instruction, {0}); + ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, values_buffer)); + const HloValue& indices_buffer = + dataflow_analysis_->GetValueDefinedAt(instruction, {1}); + ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, indices_buffer)); } else if (HloPredicateIsOp(instruction)) { Shape operand_shape = instruction->operand(0)->shape(); Shape output_shape = instruction->shape(); diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo index 3b5c92a19388d2..cac1d16a5b4b3b 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo @@ -1,17 +1,18 @@ // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s - +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: bf16[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: bf16[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = bf16[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (bf16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = bf16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = bf16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = bf16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = bf16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} ROOT convolution.3 = bf16[1,64,64,32]{3,2,1,0} convolution(Arg_0.1, Arg_1.2), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f -} +} \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo index 1b82bb55c80b2b..5d4b5daf721ee1 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo @@ -1,17 +1,18 @@ // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s - +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: f8e4m3fn[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = f8e4m3fn[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (f8e4m3fn[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f8e4m3fn[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = f8e4m3fn[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = f8e4m3fn[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} ROOT convolution.3 = f8e4m3fn[1,64,64,32]{3,2,1,0} convolution(Arg_0.1, Arg_1.2), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f -} +} \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo index d5baeb8a42af7d..ec15e5da5f3e92 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo @@ -1,17 +1,18 @@ // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/v100.txtpb --split-input-file | FileCheck %s - +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: f16[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: f16[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = f16[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (f16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = f16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = f16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = f16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} ROOT convolution.3 = f16[1,64,64,32]{3,2,1,0} convolution(Arg_0.1, Arg_1.2), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f -} +} \ No newline at end of file diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc index 909690650944b6..77fe80a539d3a0 100644 --- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc @@ -723,7 +723,8 @@ class PriorityFusionQueue { producer, consumer, use_multi_output_fusion); absl::StatusOr result_or_status = - combined_gpu_performance_model_.TryFindBestTilingForFusion(*fusion); + combined_gpu_performance_model_.TryFindBestTilingForFusionAsync(*fusion) + .Await(); // Convert absl::Status into FusionDecision. We don't distinguish between // status and FusionDecision here, because both indicate that tile analysis diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc index b3c092c3d9db6e..64e0197f544824 100644 --- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc @@ -1329,7 +1329,7 @@ TEST_P(PriorityFusionTest, DoNotFuseInsideReducer) { } TEST_P(PriorityFusionTest, SkipsTilingsWithInfiniteRuntime) { - // This test verifies the fix in TryFindBestTilingForFusion that skips + // This test verifies the fix in TryFindBestTilingForFusionAsync that skips // tilings with infinite runtime estimates. // // The fix: After estimating runtime for each tiling candidate, check if diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc index d863bea48b0f4b..cbc5735abd8af8 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc @@ -356,9 +356,10 @@ DecideIfShouldFuseAndMaybeSetBlockLevelParameters( bool use_cost_model_to_evaluate_fusions) { auto fusion_adaptor = HloFusionAdaptor::ForInstruction(normalization_fusion); - ABSL_ASSIGN_OR_RETURN( - TiledRunTimeDataOrError tiled_runtime_data_or, - indexing_performance_model.TryFindBestTilingForFusion(*fusion_adaptor)); + ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or, + indexing_performance_model + .TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); if (const auto* fusion_decision = std::get_if(&tiled_runtime_data_or)) { diff --git a/third_party/xla/xla/backends/profiler/cpu/BUILD b/third_party/xla/xla/backends/profiler/cpu/BUILD index 9487b110986d65..1b5b07f98ff20a 100644 --- a/third_party/xla/xla/backends/profiler/cpu/BUILD +++ b/third_party/xla/xla/backends/profiler/cpu/BUILD @@ -79,15 +79,17 @@ cc_library( ]), deps = [ "//xla:status_macros", - "//xla/backends/profiler/util:metadata_registry", + "//xla/backends/profiler/util:metadata_registry", # buildcleaner: keep "//xla/python/profiler/internal:python_hooks", "//xla/tsl/platform:logging", + "//xla/tsl/profiler/utils:time_utils", "//xla/tsl/profiler/utils:xplane_builder", "//xla/tsl/profiler/utils:xplane_schema", "//xla/tsl/profiler/utils:xplane_utils", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@rules_python//python/cc:current_py_cc_headers", # buildcleaner: keep "@tsl//tsl/profiler/lib:profiler_interface", "@tsl//tsl/profiler/protobuf:xplane_proto_cc", ], diff --git a/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc b/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc index f09bb2b1e33dd5..38dce9159357dd 100644 --- a/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc +++ b/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc @@ -14,8 +14,10 @@ limitations under the License. ==============================================================================*/ #include "xla/backends/profiler/cpu/python_tracer.h" +#include #include #include +#include #include #include @@ -25,6 +27,7 @@ limitations under the License. #include "xla/python/profiler/internal/python_hooks.h" #include "xla/status_macros.h" #include "xla/tsl/platform/logging.h" +#include "xla/tsl/profiler/utils/time_utils.h" #include "xla/tsl/profiler/utils/xplane_builder.h" #include "xla/tsl/profiler/utils/xplane_schema.h" #include "xla/tsl/profiler/utils/xplane_utils.h" @@ -57,6 +60,8 @@ class PythonTracer : public tsl::profiler::ProfilerInterface { private: bool recording_ = false; const PythonHooksOptions options_; + // Start timestamp (in nanoseconds) of the current recording/consume window. + uint64_t start_timestamp_ns_ = 0; std::unique_ptr context_; PythonTracer(const PythonTracer&) = delete; @@ -71,6 +76,7 @@ absl::Status PythonTracer::Start() { // TENSORFLOW_STATUS_OK } VLOG(1) << __FUNCTION__; recording_ = true; + start_timestamp_ns_ = tsl::profiler::GetCurrentTimeNanos(); PythonHooks::GetSingleton()->Start(options_); return absl::OkStatus(); } @@ -98,6 +104,8 @@ absl::Status PythonTracer::CollectData( // TENSORFLOW_STATUS_OK absl::StatusOr PythonTracer::Consume() { VLOG(1) << "PythonTracer::Consume called, recording=" << recording_; PythonTracerChunk chunk; + chunk.start_timestamp_ns = + std::exchange(start_timestamp_ns_, tsl::profiler::GetCurrentTimeNanos()); if (recording_) { chunk.consumed_data = PythonHooks::GetSingleton()->Consume(); } else if (context_) { @@ -110,7 +118,7 @@ absl::StatusOr PythonTracer::Consume() { size_t estimated_size = 0; size_t total_events = 0; - for (const auto& thread_data : chunk.consumed_data) { + for (const PerThreadConsumeData& thread_data : chunk.consumed_data) { estimated_size += sizeof(PerThreadConsumeData) + thread_data.events.size() * sizeof(TraceEventInfo); total_events += thread_data.events.size(); @@ -141,10 +149,36 @@ absl::Status PythonTracer::Serialize(std::any data, space, tsl::profiler::kPythonTracerPlaneName); tsl::profiler::XPlaneBuilder plane(raw_plane); - for (const auto& thread_data : chunk->consumed_data) { + for (const PerThreadConsumeData& thread_data : chunk->consumed_data) { + // Skip threads with no recorded events. + if (thread_data.events.empty()) { + continue; + } + + // Anchor the line to the chunk start timestamp, clamping down to the + // earliest event start timestamp if any event started before the chunk + // window (or if the chunk start timestamp was unset). This guarantees + // non-negative event offsets and prevents int64_t picosecond overflow. + uint64_t line_start_ns = chunk->start_timestamp_ns > 0 + ? chunk->start_timestamp_ns + : thread_data.events.front().start_time_ns; + for (const TraceEventInfo& event : thread_data.events) { + line_start_ns = std::min(line_start_ns, event.start_time_ns); + } + tsl::profiler::XLineBuilder line = plane.GetOrCreateLine(thread_data.thread_id); - for (const auto& event : thread_data.events) { + if (line.NumEvents() == 0 && line.TimestampNs() == 0) { + // Newly created lines have no events yet (since empty event chunks are + // skipped above), so we can set the initial baseline timestamp directly. + line.SetTimestampNs(line_start_ns); + } else if (line_start_ns < static_cast(line.TimestampNs())) { + // If this is a subsequent chunk, with event timestamps earlier than the + // existing line start timestamp, we adjust the line start timestamp and + // event offsets. + line.SetTimestampNsAndAdjustEventOffsets(line_start_ns); + } + for (const TraceEventInfo& event : thread_data.events) { tsl::profiler::XEventBuilder xevent = line.AddEvent(*plane.GetOrCreateEventMetadata(event.name)); xevent.SetTimestampNs(event.start_time_ns); diff --git a/third_party/xla/xla/backends/profiler/cpu/python_tracer.h b/third_party/xla/xla/backends/profiler/cpu/python_tracer.h index 76c4d0d498ab08..5d9b5fc8fef1be 100644 --- a/third_party/xla/xla/backends/profiler/cpu/python_tracer.h +++ b/third_party/xla/xla/backends/profiler/cpu/python_tracer.h @@ -15,17 +15,22 @@ limitations under the License. #ifndef XLA_BACKENDS_PROFILER_CPU_PYTHON_TRACER_H_ #define XLA_BACKENDS_PROFILER_CPU_PYTHON_TRACER_H_ +#include #include #include #include "xla/python/profiler/internal/python_hooks.h" #include "tsl/profiler/lib/profiler_interface.h" -#include "tsl/profiler/protobuf/xplane.pb.h" namespace xla { namespace profiler { +// A chunk of Python trace data collected during a streaming profiling interval. struct PythonTracerChunk { + // Timestamp (in nanoseconds) when this chunk's profiling interval began + // (either at Start() or at the preceding Consume() call). Used as the + // initial baseline timestamp for thread XLines during Serialize(). + uint64_t start_timestamp_ns = 0; std::vector consumed_data; }; diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc index 8b2abc3b0c8219..175997a73cfe58 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc @@ -1515,6 +1515,65 @@ absl::StatusOr PropagateTileToOutputForBitcastOp( return Tiles{std::move(output_tile)}; } +// Calculates the required input tile for a ReduceScatter operation. +// +// In a ReduceScatter, the input tensor is divided into "shards" along the +// scatter dimension and distributed across R replicas. For example: +// * Replicas (R) = 2 +// * Input shape = [128, 256] +// * Output shape = [64, 256] (Each replica gets a 64-row shard) +// +// This method does two things to the tile along the scattered dimension: +// +// 1. Scales it up by R: +// Because the data will be scattered, a thread block needs to process R +// times as much input data. If the block is assigned an output tile of +// 8 rows, it must actually compute 16 rows of the input. We multiply +// the offset, size, and upper bounds by R to reflect this. +// +// 2. Enforces single-destination routing: +// We add a divisibility constraint to ensure that our scaled-up input +// tile (e.g., 16 rows) fits perfectly inside a single replica's shard +// (64 rows). This guarantees that the tile never crosses the boundary +// between shards (e.g., crossing row 64), meaning this thread block +// only ever has to send its data to exactly *one* target GPU. +absl::StatusOr PropagateTileToInputForReduceScatterOp( + const TilingSpace& tiling_space, + const HloReduceScatterInstruction& reduce_scatter, + const Tile& output_tile) { + // Crash OK. Must be checked while forming the fusion. + CHECK_EQ(reduce_scatter.operand_count(), 1) + << "Multi-operand ReduceScatter is not yet supported."; + const int64_t scatter_dim = reduce_scatter.scatter_dimension(); + const int64_t shard_size = reduce_scatter.shape().dimensions(scatter_dim); + const int64_t num_replicas = + reduce_scatter.operand(0)->shape().dimensions(scatter_dim) / shard_size; + + SmallVector dim_tiles(output_tile.dim_tiles()); + DimTile& scatter_tile = dim_tiles[scatter_dim]; + if (!IsConstantValue(scatter_tile.stride, 1)) { + return absl::UnimplementedError( + absl::StrCat("Non-unit stride on the scatter dimension of ", + reduce_scatter.ToString())); + } + scatter_tile.offset = scatter_tile.offset * num_replicas; + scatter_tile.size = scatter_tile.size * num_replicas; + scatter_tile.upper_bound = scatter_tile.upper_bound * num_replicas; + + if (!tiling_space.IsSymbolic()) { + std::optional size = TryGetConstantValue(scatter_tile.size); + if (!size.has_value() || shard_size % *size != 0 || + !scatter_tile.offset.Canonicalize().IsMultipleOf(*size)) { + return absl::FailedPreconditionError(absl::StrCat( + "Tiling propagation rejected for ", reduce_scatter.ToString(), + ": operand tile ", scatter_tile.size.ToString(), " at ", + scatter_tile.offset.ToString(), + " crosses shard boundary (shard_size=", shard_size, ")")); + } + } + return Tiles{output_tile.CloneWithNewDims(std::move(dim_tiles))}; +} + } // namespace std::string ToString(const Tiles& tiles) { @@ -1589,6 +1648,10 @@ absl::StatusOr PropagateTileToInput(TilingSpace& tiling_space, if (hlo.opcode() == HloOpcode::kAllGather) { return PropagateTileToInputForAllGatherOp(tiling_space, hlo, output_tile); } + if (hlo.opcode() == HloOpcode::kReduceScatter) { + return PropagateTileToInputForReduceScatterOp( + tiling_space, *Cast(&hlo), output_tile); + } if (hlo.opcode() == HloOpcode::kBitcast) { return PropagateTileToInputForBitcastOp(hlo, output_tile); } diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc index 070653530648b2..3182cd63ce52f8 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc @@ -198,6 +198,161 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfAllGatherOp) { )")); } +TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOp) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[128,256] parameter(0) + ROOT reduce_scatter = f32[64,256] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + ASSERT_OK_AND_ASSIGN( + auto tiled_operands, + PropagateTileToInput( + *tiling_space, *root, + GetTestTile(*tiling_space, root->shape().dimensions()), 0)); + EXPECT_THAT(tiled_operands, MatchToString(R"( + 0) (tid_0, tid_1) + -> offsets [tid_0 * ts_0 * 2, tid_1 * ts_1] + sizes [ts_0 * 2, ts_1] + strides [1, 2] + upper bounds [128, 256] + )")); +} + +TEST_F(TilePropagationTest, CanPropagateFromReduceScatterToDotOp) { + HloInstruction* fusion = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + fusion { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + dot = f32[128,512] dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT reduce_scatter = f32[64,512] reduce-scatter(dot), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + ENTRY e { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + ROOT fusion = f32[64,512] fusion(p0, p1), kind=kCustom, calls=fusion + } + )"); + const HloInstruction* root = fusion->fused_expression_root(); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(fusion), + &mlir_context_)); + ASSERT_OK_AND_ASSIGN( + auto rs_operands, + PropagateTileToInput( + *tiling_space, *root, + GetTestTile(*tiling_space, root->shape().dimensions()), 0)); + ASSERT_OK_AND_ASSIGN(auto dot_operands, + PropagateTileToInput(*tiling_space, *root->operand(0), + rs_operands[0], 0)); + EXPECT_THAT(dot_operands, MatchToString(R"( + 0) (tid_0, tid_1, tid_2) + -> offsets [tid_0 * ts_0 * 2, tid_2 * ts_2] + sizes [ts_0 * 2, ts_2] + strides [1, 1] + upper bounds [128, 256] + 1) (tid_0, tid_1, tid_2) + -> offsets [tid_2 * ts_2, tid_1 * ts_1] + sizes [ts_2, ts_1] + strides [1, 2] + upper bounds [256, 512] + )")); +} + +TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOpScatterDim1) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[128,256] parameter(0) + ROOT reduce_scatter = f32[128,128] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={1}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + ASSERT_OK_AND_ASSIGN(auto tiled_operands, + PropagateTileToInput(*tiling_space, *root, + tiling_space->tiled_roots()[0], 0)); + EXPECT_THAT(tiled_operands, MatchToString(R"( + 0) (tid_0, tid_1) + -> offsets [tid_0 * ts_0, tid_1 * ts_1 * 2] + sizes [ts_0, ts_1 * 2] + strides [1, 1] + upper bounds [128, 256] + )")); +} + +TEST_F(TilePropagationTest, ReduceScatterRejectsNonUnitScatterStride) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[128,256] parameter(0) + ROOT reduce_scatter = f32[128,128] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={1}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + // GetTestTile assigns stride = index + 1 = 2 on dimension 1. + EXPECT_THAT(PropagateTileToInput( + *tiling_space, *root, + GetTestTile(*tiling_space, root->shape().dimensions()), 0), + StatusIs(absl::StatusCode::kUnimplemented)); +} + +TEST_F(TilePropagationTest, ReduceScatterRejectsTileCrossingShards) { + HloInstruction* root = ParseAndGetRoot(R"( + HloModule m + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + ENTRY e { + p0 = f32[192,256] parameter(0) + ROOT reduce_scatter = f32[96,256] reduce-scatter(p0), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + )"); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr tiling_space, + TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root), + &mlir_context_)); + // Operand tile 32 * 2 = 64: block 1 reads rows [64, 128), crossing 96. + ASSERT_OK(tiling_space->AssignTileSizes({32, 256})); + EXPECT_THAT(PropagateTileToInput(*tiling_space, *root, + tiling_space->tiled_roots()[0], 0), + StatusIs(absl::StatusCode::kFailedPrecondition)); +} + TEST_F(TilePropagationTest, CanPropagateToInputOfBroadcastOp) { HloInstruction* root = ParseAndGetRoot(R"( HloModule m diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc index 8e1cb53ad9576e..e9fb30a0260fce 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc @@ -700,6 +700,47 @@ TEST_P(TileAnalysisTest, CollectiveDotBasic) { )")); } +TEST_P(TileAnalysisTest, CollectiveDotReduceScatter) { + ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation, + ParseAndTile(R"hlo( + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + + fusion { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + dot = f32[128,512] dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT rs = f32[64,512] reduce-scatter(dot), replica_groups={{0,1}}, dimensions={0}, to_apply=add + } + + ENTRY main { + p0 = f32[128,256] parameter(0) + p1 = f32[256,512] parameter(1) + ROOT fusion = f32[64,512] fusion(p0, p1), kind=kCustom, calls=fusion + })hlo", + {8, 32, 32})); + + EXPECT_THAT(tiled_computation, MatchString(R"( + Dimensions: + 0 type: parallel size: 64 tile size: 8 dim ID:0 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add + 1 type: parallel size: 512 tile size: 32 dim ID:1 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add + 2 type: sequential size: 256 tile size: 32 dim ID:2 hlo: %dot = f32[128,512]{1,0} dot(%p0, %p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + Root tiles: + 0 root tile: offsets [tid_0 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512] + + Tiled HLO: + dot.tile_0 = dot(p0.1.tile_0, p1.1.tile_0) offsets [tid_0 * 16, tid_1 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 512] + region #0 { + p0.1.tile_0 = parameter(0) offsets [tid_0 * 16, tid_2 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 256] + p1.1.tile_0 = parameter(1) offsets [tid_2 * 32, tid_1 * 32] sizes [32, 32] strides [1, 1] upper bounds [256, 512] + } + rs.tile_0 = reduce-scatter(dot.tile_0) offsets [tid_0 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512] + )")); +} + TEST_F(SingleOutputFusionTileAnalysisTest, SimpleFusion) { absl::StatusOr parse_result = ParseAndTile(R"hlo( fusion { diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc index f3494bbe77856e..f161de27bae98a 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc +++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc @@ -133,6 +133,52 @@ ArrayRef MakeArrayRef(const absl::Span span) { return ArrayRef(span.data(), span.size()); } +absl::StatusOr EmitAllGather( + EmitterContext& emitter_ctx, const HloAllGatherInstruction* all_gather, + const ge::TiledHloInstruction& tiled_all_gather, ValueRange operands) { + if (all_gather->device_list()->replica_groups().empty()) { + return Internal( + "Triton emitting AllGather without replica groups is not supported."); + } + + llvm::SmallVector flattened_replica_group_ids; + for (const auto& replica_group : all_gather->replica_groups()) { + for (const auto& replica_id : replica_group.replica_ids()) { + flattened_replica_group_ids.push_back(replica_id); + } + } + + std::optional channel_handle = all_gather->channel_id(); + bool use_global_device_ids = all_gather->use_global_device_ids(); + + ImplicitLocOpBuilder& b = emitter_ctx.b(); + ABSL_ASSIGN_OR_RETURN( + auto output_element_type, + xtile::PrimitiveTypeToMlirType(b, all_gather->shape().element_type())); + ABSL_ASSIGN_OR_RETURN(SmallVector tile_sizes, + tiled_all_gather.tile().GetStaticTileSizes()); + auto output_type = + mlir::RankedTensorType::get(tile_sizes, output_element_type); + + auto replica_groups_type = mlir::RankedTensorType::get( + {static_cast(all_gather->replica_groups().size()), + static_cast( + all_gather->replica_groups()[0].replica_ids_size())}, + b.getI64Type()); + auto replica_groups_attr = mlir::DenseIntElementsAttr::get( + replica_groups_type, flattened_replica_group_ids); + auto channel_handle_attr = + channel_handle ? mlir::stablehlo::ChannelHandleAttr::get(b.getContext(), + *channel_handle, + /*type=*/0) + : nullptr; + auto all_gather_op = mlir::stablehlo::AllGatherOp::create( + b, mlir::TypeRange{output_type}, operands, + all_gather->all_gather_dimension(), replica_groups_attr, + channel_handle_attr, use_global_device_ids); + return mlir::cast(all_gather_op.getResult(0)); +} + absl::StatusOr EmitAllReduce( EmitterContext& emitter_ctx, const HloAllReduceInstruction* all_reduce, const ge::TiledHloInstruction& tiled_all_reduce, ValueRange operands) { @@ -1884,8 +1930,8 @@ absl::StatusOr EmitTiledHloInstruction( // Please keep the cases in alphabetical order. switch (hlo->opcode()) { case HloOpcode::kAllGather: { - // AllGather is a no-op. Tile extraction handles the data movement. - return emitter_ctx.TiledHloToTensorValue(*tiled_hlo.operand(0)); + return EmitAllGather(emitter_ctx, xla::Cast(hlo), + tiled_hlo, operands); } case HloOpcode::kAllReduce: { const HloComputation* computation = diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc index 331bb845448054..91a224866506bb 100644 --- a/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc +++ b/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc @@ -84,18 +84,18 @@ std::optional IsLegalStablehloOp(mlir::Operation* op) { if (mlir::isa< // go/keep-sorted start mlir::stablehlo::AbsOp, mlir::stablehlo::AddOp, - mlir::stablehlo::AllReduceOp, mlir::stablehlo::AndOp, - mlir::stablehlo::BroadcastInDimOp, mlir::stablehlo::CompareOp, - mlir::stablehlo::ComplexOp, mlir::stablehlo::ConstantOp, - mlir::stablehlo::ConvertOp, mlir::stablehlo::DivOp, - mlir::stablehlo::DotGeneralOp, mlir::stablehlo::ImagOp, - mlir::stablehlo::MaxOp, mlir::stablehlo::MinOp, - mlir::stablehlo::MulOp, mlir::stablehlo::NegOp, mlir::stablehlo::OrOp, - mlir::stablehlo::PowOp, mlir::stablehlo::RealOp, - mlir::stablehlo::ReduceOp, mlir::stablehlo::RemOp, - mlir::stablehlo::ReshapeOp, mlir::stablehlo::ReturnOp, - mlir::stablehlo::SubtractOp, mlir::stablehlo::TransposeOp, - mlir::stablehlo::XorOp + mlir::stablehlo::AllGatherOp, mlir::stablehlo::AllReduceOp, + mlir::stablehlo::AndOp, mlir::stablehlo::BroadcastInDimOp, + mlir::stablehlo::CompareOp, mlir::stablehlo::ComplexOp, + mlir::stablehlo::ConstantOp, mlir::stablehlo::ConvertOp, + mlir::stablehlo::DivOp, mlir::stablehlo::DotGeneralOp, + mlir::stablehlo::ImagOp, mlir::stablehlo::MaxOp, + mlir::stablehlo::MinOp, mlir::stablehlo::MulOp, + mlir::stablehlo::NegOp, mlir::stablehlo::OrOp, mlir::stablehlo::PowOp, + mlir::stablehlo::RealOp, mlir::stablehlo::ReduceOp, + mlir::stablehlo::RemOp, mlir::stablehlo::ReshapeOp, + mlir::stablehlo::ReturnOp, mlir::stablehlo::SubtractOp, + mlir::stablehlo::TransposeOp, mlir::stablehlo::XorOp // go/keep-sorted end >(op)) { return std::nullopt; diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td index 3f49d7032eba1f..8a9d655d3b71cc 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td @@ -34,6 +34,7 @@ def XTileDialect : Dialect { ]; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; } #endif // XLA_CODEGEN_XTILE_IR_XTILE_DIALECT diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 4a5a969daac4eb..a86920cee08c4d 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -303,6 +303,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_trace_annotation_level(0); opts.set_xla_gpu_enable_cupti_multi_subscriber(true); + opts.add_xla_gpu_enable_command_buffer(DebugOptions::COLLECTIVES_KERNEL); opts.add_xla_gpu_enable_command_buffer(DebugOptions::CONDITIONAL); opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUBLAS); opts.add_xla_gpu_enable_command_buffer(DebugOptions::CUBLASLT); @@ -522,7 +523,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_experimental_matmul_perf_table_path(""); // TODO(b/366475196): Create XLA GPU without cuDNN, cuBLAS. opts.set_xla_gpu_experimental_disable_binary_libraries(false); - opts.set_xla_gpu_experimental_enable_conv_fusion(false); + opts.set_xla_gpu_experimental_enable_conv_fusion(true); opts.set_xla_gpu_dot_merger_threshold_mb(64); opts.set_xla_enable_fast_math(false); opts.set_xla_gpu_experimental_parallel_collective_overlap_limit(1); diff --git a/third_party/xla/xla/debug_options_parsers_test.cc b/third_party/xla/xla/debug_options_parsers_test.cc index 78d35a6bd31642..8c08a475fb0dff 100644 --- a/third_party/xla/xla/debug_options_parsers_test.cc +++ b/third_party/xla/xla/debug_options_parsers_test.cc @@ -442,15 +442,15 @@ TEST(ParseRepeatedEnumModifiersTest, Invalid) { TEST(ParseRepeatedEnumFlagsTest, CommandBufferCmdType) { DebugOptions debug_options = DefaultDebugOptionsIgnoringFlags(); - // Check that the default setting has 6 types. + // Check that the default setting has 8 types. const auto& enabled_types = debug_options.xla_gpu_enable_command_buffer(); - ASSERT_EQ(enabled_types.size(), 7); + ASSERT_EQ(enabled_types.size(), 8); ASSERT_THAT( enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLAS, - DebugOptions::CUBLASLT, DebugOptions::CUDNN, - DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION, - DebugOptions::FUSION)); + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, DebugOptions::CONDITIONAL, + DebugOptions::CUBLAS, DebugOptions::CUBLASLT, + DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, + DebugOptions::DYNAMIC_SLICE_FUSION, DebugOptions::FUSION)); // Initialize the flag objects. std::vector flag_objects; @@ -459,30 +459,33 @@ TEST(ParseRepeatedEnumFlagsTest, CommandBufferCmdType) { // Removing options from the existing setting. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=-fusion,-cublas"); ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", flag_objects); - EXPECT_EQ(enabled_types.size(), 5); + EXPECT_EQ(enabled_types.size(), 6); EXPECT_THAT(enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, + DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION)); // Removing an option that isn't there and adding a duplicate. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=+cublaslt,-fusion"); ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", flag_objects); - EXPECT_EQ(enabled_types.size(), 5); + EXPECT_EQ(enabled_types.size(), 6); EXPECT_THAT(enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, + DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION)); // Adding an option. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=+cublas"); ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", flag_objects); - EXPECT_EQ(enabled_types.size(), 6); + EXPECT_EQ(enabled_types.size(), 7); EXPECT_THAT( enabled_types, - ElementsAre(DebugOptions::CONDITIONAL, DebugOptions::CUBLASLT, - DebugOptions::CUDNN, DebugOptions::CUSTOM_CALL, - DebugOptions::DYNAMIC_SLICE_FUSION, DebugOptions::CUBLAS)); + ElementsAre(DebugOptions::COLLECTIVES_KERNEL, DebugOptions::CONDITIONAL, + DebugOptions::CUBLASLT, DebugOptions::CUDNN, + DebugOptions::CUSTOM_CALL, DebugOptions::DYNAMIC_SLICE_FUSION, + DebugOptions::CUBLAS)); // Overwriting the default setting. SetXlaFlagsEnvVar("--xla_gpu_enable_command_buffer=custom_call,fusion"); diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc index 83ca7aa091d94c..76d6ecf2ab1a3c 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc @@ -107,12 +107,13 @@ using absl::StrCat; HloDataflowAnalysis::HloDataflowAnalysis( const HloModule& module, bool ssa_form, bool bitcast_defines_value, absl::flat_hash_set execution_threads, - bool propagate_through_calls) + bool propagate_through_calls, bool propagate_through_control_flow) : module_(module), execution_threads_(std::move(execution_threads)), ssa_form_(ssa_form), bitcast_defines_value_(bitcast_defines_value), propagate_through_calls_(propagate_through_calls), + propagate_through_control_flow_(propagate_through_control_flow), call_graph_(CallGraph::Build(&module)) {} bool HloDataflowAnalysis::AreTransitiveUsesElementwiseOrTuple( @@ -477,8 +478,10 @@ bool HloDataflowAnalysis::UpdateAsyncChainOutputValueSet( HloInstruction* async_op) { CHECK(async_op->IsAsynchronous()); bool changed = false; - bool is_thread_included = HloInstruction::IsThreadIncluded( - async_op->async_execution_thread(), execution_threads_); + bool is_thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded(async_op->async_execution_thread(), + execution_threads_); if (!is_thread_included && async_op->opcode() == HloOpcode::kAsyncStart) { // AsyncStart in a non-included thread has the output values defined, no @@ -563,6 +566,9 @@ bool HloDataflowAnalysis::UpdateAsyncStartValueSet( changed |= UpdateAsyncChainOperandValueSet(async_start, i, async_start->operand(i)); } + if (!propagate_through_control_flow_) { + return changed; + } bool is_dus = async_start->async_wrapped_opcode() == HloOpcode::kDynamicUpdateSlice; @@ -745,6 +751,9 @@ bool HloDataflowAnalysis::UpdateCallValueSet(HloInstruction* call) { bool HloDataflowAnalysis::UpdateConditionalValueSet( HloInstruction* conditional) { + if (!propagate_through_control_flow_) { + return false; + } CHECK_EQ(conditional->opcode(), HloOpcode::kConditional); std::vector inputs(conditional->branch_count()); for (int j = 0; j < conditional->branch_count(); ++j) { @@ -784,7 +793,8 @@ bool HloDataflowAnalysis::UpdateOptimizationBarrierValueSet( // Optimization Barriers just forward their operand. Given that barriers can // have a tuple operand, we iterate through its indexes, like for copies. // Unlike copies though we also propagate the top-level value. - CHECK_EQ(barrier->opcode(), HloOpcode::kOptimizationBarrier); + CHECK(barrier->opcode() == HloOpcode::kOptimizationBarrier || + barrier->IsCustomCall(kCallMarkerAfterTarget)); bool changed = false; for (auto& pair : GetInstructionValueSet(barrier)) { const ShapeIndex& index = pair.first; @@ -884,6 +894,9 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { callsite.instruction()->operand(parameter->parameter_number()))); } } else if (opcode == HloOpcode::kWhile) { + if (!propagate_through_control_flow_) { + continue; + } // In a while instruction, the while operand (ie, the init value) and the // backedge are dataflow inputs to the parameter instruction. This is the // case for parameters of both the body and condition computations. @@ -900,6 +913,9 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { } need_phi = true; } else if (opcode == HloOpcode::kConditional) { + if (!propagate_through_control_flow_) { + continue; + } CHECK_EQ(parameter->parameter_number(), 0); auto conditional = callsite.instruction(); // Conditional has branch_count+1 operands. Operand 0 is the branch_index, @@ -921,6 +937,9 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { CHECK(found_parent); need_phi = true; } else if (opcode == HloOpcode::kAsyncStart) { + if (!propagate_through_control_flow_) { + continue; + } const HloInstruction* async_done = callsite.instruction()->async_chain_done(); // When an async chain crosses a while loop boundary, async_chain_done() @@ -977,7 +996,8 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { } bool HloDataflowAnalysis::UpdateTupleValueSet(HloInstruction* tuple) { - CHECK_EQ(tuple->opcode(), HloOpcode::kTuple); + CHECK(tuple->opcode() == HloOpcode::kTuple || + tuple->IsCustomCall(kCallMarkerBeforeTarget)); bool changed = false; for (int64_t i = 0; i < tuple->operands().size(); ++i) { // Copy the value set(s) of each operand into the respective position in the @@ -1002,6 +1022,9 @@ bool HloDataflowAnalysis::UpdateTupleValueSet(HloInstruction* tuple) { } bool HloDataflowAnalysis::UpdateWhileValueSet(HloInstruction* xla_while) { + if (!propagate_through_control_flow_) { + return false; + } CHECK_EQ(xla_while->opcode(), HloOpcode::kWhile); const InstructionValueSet* const inputs[] = { &GetInstructionValueSet(xla_while->while_body()->root_instruction()), @@ -1220,6 +1243,14 @@ bool HloDataflowAnalysis::UpdateInstructionValueSet( return UpdateCollectivePermuteDoneValueSet(instruction); case HloOpcode::kOptimizationBarrier: return UpdateOptimizationBarrierValueSet(instruction); + case HloOpcode::kCustomCall: + if (instruction->custom_call_target() == kCallMarkerBeforeTarget) { + return UpdateTupleValueSet(instruction); + } + if (instruction->custom_call_target() == kCallMarkerAfterTarget) { + return UpdateOptimizationBarrierValueSet(instruction); + } + return false; default: break; } @@ -1340,6 +1371,10 @@ void HloDataflowAnalysis::Propagate() { callsite.instruction()->opcode() == HloOpcode::kCall) { continue; } + if (!propagate_through_control_flow_ && + callsite.instruction()->opcode() != HloOpcode::kCall) { + continue; + } if (callsite.instruction()->opcode() == HloOpcode::kWhile) { // Add the while itself, and the body and condition parameters. add_to_worklist(callsite.instruction()); @@ -1395,6 +1430,9 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { const CallGraphNode& call_graph_node = call_graph_->GetNode(computation); const bool is_regular_call_computation = IsRegularCallComputation(call_graph_node); + const bool is_control_flow_computation = + !call_graph_node.caller_callsites().empty() && + !is_regular_call_computation; for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { // Create an empty shape tree. @@ -1436,9 +1474,13 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { define_all_values(); } break; - case HloOpcode::kAddDependency: case HloOpcode::kWhile: case HloOpcode::kConditional: + if (!propagate_through_control_flow_) { + define_all_values(); + } + break; + case HloOpcode::kAddDependency: case HloOpcode::kGetTupleElement: case HloOpcode::kDomain: case HloOpcode::kOptimizationBarrier: @@ -1459,7 +1501,9 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { } if (call_graph_node.caller_callsites().empty() || call_graph_node.context() == CallContext::kEmbedded || - (!propagate_through_calls_ && is_regular_call_computation)) { + (!propagate_through_calls_ && is_regular_call_computation) || + (!propagate_through_control_flow_ && + is_control_flow_computation)) { // Parameters of computations called in a parallel context (eg, map // and reduce) as well as parameters of dead computations define all // values in their output. Otherwise the values of the parameter @@ -1474,6 +1518,14 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { // values flow from their operands. define_value_at(/*index=*/{}); break; + case HloOpcode::kCustomCall: + if (instruction->custom_call_target() == kCallMarkerBeforeTarget) { + define_value_at(/*index=*/{}); + } else if (instruction->custom_call_target() != + kCallMarkerAfterTarget) { + define_all_values(); + } + break; case HloOpcode::kAsyncStart: { // AsyncStart produces a tuple of {{aliased operands}, {destination}, // contexts}. It defines all of the tuple-shaped values and the @@ -1481,8 +1533,10 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { // // If the thread is excluded, then we don't track the contained // dataflow, and define the destination values too. - bool thread_included = HloInstruction::IsThreadIncluded( - instruction->async_execution_thread(), execution_threads_); + bool thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded( + instruction->async_execution_thread(), execution_threads_); define_all_values([&](const ShapeIndex& index) { return ShapeUtil::GetSubshape(instruction->shape(), index) .IsTuple() || @@ -1493,23 +1547,55 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() { } case HloOpcode::kAsyncUpdate: { // AsyncUpdate produces a tuple of {{aliased operands}, {destination}, - // contexts}. It defines all of the tuple-shaped values and the - // contexts. + // contexts}. It defines all of the tuple values and the contexts. + // When the async thread is excluded (or control flow propagation is + // disabled) and the output at {1, ...} is late bound, AsyncUpdate + // also defines the newly bound output values. + bool thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded( + instruction->async_execution_thread(), execution_threads_); + const Shape& prev_shape = instruction->operand(0)->shape(); define_all_values([&](const ShapeIndex& index) { - return ShapeUtil::GetSubshape(instruction->shape(), index) - .IsTuple() || - index.front() > 1; + if (ShapeUtil::GetSubshape(instruction->shape(), index).IsTuple() || + index.front() > 1) { + return true; + } + if (!thread_included && index.front() == 1) { + return !ShapeUtil::IndexIsValid(prev_shape, index) || + !ShapeUtil::Compatible( + ShapeUtil::GetSubshape(instruction->shape(), index), + ShapeUtil::GetSubshape(prev_shape, index)); + } + return false; }); break; } - case HloOpcode::kAsyncDone: + case HloOpcode::kAsyncDone: { // AsyncDone's output aliases its output. It defines all remaining - // tuple-shaped values. + // tuple values, plus any late bound output values when the async + // thread is excluded (or control flow propagation is disabled). + bool thread_included = + propagate_through_control_flow_ && + HloInstruction::IsThreadIncluded( + instruction->async_execution_thread(), execution_threads_); + const Shape& prev_shape = instruction->operand(0)->shape(); define_all_values([&](const ShapeIndex& index) { - return ShapeUtil::GetSubshape(instruction->shape(), index) - .IsTuple(); + if (ShapeUtil::GetSubshape(instruction->shape(), index).IsTuple()) { + return true; + } + if (!thread_included) { + ShapeIndex src_index = {1}; + src_index.insert(src_index.end(), index.begin(), index.end()); + return !ShapeUtil::IndexIsValid(prev_shape, src_index) || + !ShapeUtil::Compatible( + ShapeUtil::GetSubshape(instruction->shape(), index), + ShapeUtil::GetSubshape(prev_shape, src_index)); + } + return false; }); break; + } case HloOpcode::kCopyStart: // CopyStart produces a tuple of {destination buffer, aliased operand, // U32 context}. @@ -1655,13 +1741,14 @@ absl::StatusOr> HloDataflowAnalysis::Run( const HloModule& module, bool ssa_form, bool bitcast_defines_value, absl::flat_hash_set execution_threads, bool propagate_through_calls, - std::optional> precompute_uses) { + std::optional> precompute_uses, + bool propagate_through_control_flow) { VLOG(1) << "HloDataflowAnalysis::Run on module " << module.name(); XLA_VLOG_LINES(2, module.ToString()); - auto dataflow_analysis = absl::WrapUnique( - new HloDataflowAnalysis(module, ssa_form, bitcast_defines_value, - execution_threads, propagate_through_calls)); + auto dataflow_analysis = absl::WrapUnique(new HloDataflowAnalysis( + module, ssa_form, bitcast_defines_value, execution_threads, + propagate_through_calls, propagate_through_control_flow)); ABSL_RETURN_IF_ERROR(dataflow_analysis->RunImpl()); if (precompute_uses.has_value()) { std::vector values; diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h index 202a9c6cb60cd9..b05ca2ca8d0447 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h +++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.h @@ -88,7 +88,8 @@ class HloDataflowAnalysis { absl::flat_hash_set execution_threads = {}, bool propagate_through_calls = true, std::optional> precompute_uses = - std::nullopt); + std::nullopt, + bool propagate_through_control_flow = true); // Returns true if 'instruction' defines an HLO value at the given shape index // of its output. @@ -219,7 +220,8 @@ class HloDataflowAnalysis { HloDataflowAnalysis(const HloModule& module, bool ssa_form, bool bitcast_defines_value, absl::flat_hash_set execution_threads, - bool propagate_through_calls = true); + bool propagate_through_calls = true, + bool propagate_through_control_flow = true); // Runs dataflow analysis on the module attached to this HloDataflowAnalysis. absl::Status RunImpl(); @@ -335,6 +337,7 @@ class HloDataflowAnalysis { const bool ssa_form_; const bool bitcast_defines_value_; bool propagate_through_calls_ = true; + bool propagate_through_control_flow_ = true; std::unique_ptr call_graph_; diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc index 9265fb58d1f346..fbd952fbe8fe62 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc @@ -4332,5 +4332,109 @@ ENTRY main { } } +TEST_F(HloDataflowAnalysisTest, DisablePropagateThroughControlFlow) { + const char* hlo_text = R"hlo( +HloModule module + +callee { + ROOT call_param = f32[] parameter(0) +} + +while_cond { + param = f32[] parameter(0) + ROOT cond = pred[] constant(false) +} + +while_body { + param = f32[] parameter(0) + ROOT add = f32[] add(param, param) +} + +ENTRY main { + const0 = f32[] constant(1.0) + call_inst = f32[] call(const0), to_apply=callee + ROOT while_inst = f32[] while(call_inst), condition=while_cond, body=while_body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule( + hlo_text, GetModuleConfigForTest())); + HloInstruction* const0 = FindInstruction(module_.get(), "const0"); + HloInstruction* call_inst = FindInstruction(module_.get(), "call_inst"); + HloInstruction* call_param = + module_->GetComputationWithName("callee")->parameter_instruction(0); + HloInstruction* while_inst = FindInstruction(module_.get(), "while_inst"); + HloInstruction* body_param = + module_->GetComputationWithName("while_body")->parameter_instruction(0); + + { + ASSERT_OK_AND_ASSIGN( + auto analysis, + HloDataflowAnalysis::Run( + *module_, /*ssa_form=*/false, /*bitcast_defines_value=*/false, + /*execution_threads=*/{}, /*propagate_through_calls=*/false, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(call_inst)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(call_param)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(while_inst)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(body_param)); + EXPECT_THAT(analysis->GetValueSet(body_param).values(), + ::testing::Not( + ::testing::Contains(&analysis->GetValueDefinedAt(const0)))); + } + { + ASSERT_OK_AND_ASSIGN( + auto analysis, + HloDataflowAnalysis::Run( + *module_, /*ssa_form=*/false, /*bitcast_defines_value=*/false, + /*execution_threads=*/{}, /*propagate_through_calls=*/true, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); + EXPECT_FALSE(analysis->ValueIsDefinedAt(call_inst)); + EXPECT_FALSE(analysis->ValueIsDefinedAt(call_param)); + EXPECT_THAT(analysis->GetValueSet(call_inst).values(), + ::testing::Contains(&analysis->GetValueDefinedAt(const0))); + EXPECT_TRUE(analysis->ValueIsDefinedAt(while_inst)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(body_param)); + EXPECT_THAT(analysis->GetValueSet(body_param).values(), + ::testing::Not( + ::testing::Contains(&analysis->GetValueDefinedAt(const0)))); + } +} + +TEST_F(HloDataflowAnalysisTest, CallMarkerCustomCalls) { + const char* hlo_text = R"hlo( +HloModule module + +ENTRY main { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + before = (f32[4], f32[4]) custom-call(p0, p1), custom_call_target="__xla_internal_call_marker_before" + gte0 = f32[4] get-tuple-element(before), index=0 + gte1 = f32[4] get-tuple-element(before), index=1 + add = f32[4] add(gte0, gte1) + ROOT after = f32[4] custom-call(add), custom_call_target="__xla_internal_call_marker_after" +} +)hlo"; + ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule( + hlo_text, GetModuleConfigForTest())); + HloInstruction* p0 = FindInstruction(module_.get(), "p0"); + HloInstruction* p1 = FindInstruction(module_.get(), "p1"); + HloInstruction* before = FindInstruction(module_.get(), "before"); + HloInstruction* add = FindInstruction(module_.get(), "add"); + HloInstruction* after = FindInstruction(module_.get(), "after"); + + ASSERT_OK_AND_ASSIGN(auto analysis, HloDataflowAnalysis::Run(*module_)); + EXPECT_TRUE(analysis->ValueIsDefinedAt(before, {})); + EXPECT_FALSE(analysis->ValueIsDefinedAt(before, {0})); + EXPECT_EQ(&analysis->GetUniqueValueAt(before, {0}), + &analysis->GetValueDefinedAt(p0)); + EXPECT_EQ(&analysis->GetUniqueValueAt(before, {1}), + &analysis->GetValueDefinedAt(p1)); + EXPECT_FALSE(analysis->ValueIsDefinedAt(after)); + EXPECT_EQ(&analysis->GetUniqueValueAt(after), + &analysis->GetValueDefinedAt(add)); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index d1e46162ccd829..9b316fb58ed935 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -46,7 +46,9 @@ limitations under the License. #include "xla/layout.h" #include "xla/literal.h" #include "xla/literal_pool.h" +#include "xla/primitive_util.h" #include "xla/printer.h" +#include "xla/protobuf_util.h" #include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -81,6 +83,10 @@ class HloDimensionsInstruction : public HloInstruction { } } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), dimensions()); + } + protected: HloDimensionsInstruction(HloOpcode opcode, const Shape& shape, absl::Span dimensions) @@ -462,6 +468,10 @@ class HloCompareInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kCompare; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), direction(), order()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -1081,6 +1091,10 @@ class HloCollectivePermuteInstruction : public HloChannelInstruction { hlo->opcode() == HloOpcode::kCollectivePermuteStart; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), source_target_pairs()); + } + // Whether this is an in-place collective permute (with dynamic slice // operands). Derived from the presence of slice_sizes. bool inplace() const { return !slice_sizes_.empty(); } @@ -1517,6 +1531,11 @@ class HloSliceInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kSlice; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), slice_starts(), slice_limits(), + slice_strides()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -1579,11 +1598,19 @@ class HloConstantInstruction : public HloInstruction { return false; } - // Add literal to the hash state. + // Hashes the literal the way Identical compares it: by value, ignoring the + // literal's layout, over at most 64 bytes. Literal equality compares sub + // byte types under a bit mask and dynamic arrays only up to their dynamic + // sizes, so their bytes stay out. void HashAdditionalAttributes(absl::HashState h) const override { - if (HasLiteral()) { - absl::HashState::combine(std::move(h), - Literal::AbslHashable(literal())); + if (!HasLiteral()) { + return; + } + const Shape& shape = literal().shape(); + if (shape.IsArray() && shape.is_static() && + !primitive_util::IsSubByteNonPredType(shape.element_type())) { + LiteralBase::Hash(std::move(h), literal()); } } @@ -1976,6 +2003,10 @@ class HloGetTupleElementInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kGetTupleElement; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), tuple_index()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2167,6 +2198,11 @@ class HloConvolutionInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kConvolution; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), protobuf_util::ProtobufHashBySerialization(window())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2270,6 +2306,11 @@ class HloReduceWindowInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kReduceWindow; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), protobuf_util::ProtobufHashBySerialization(window())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2461,6 +2502,10 @@ class HloCustomCallInstruction : public HloCallableInstruction { return hlo->opcode() == HloOpcode::kCustomCall; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), custom_call_target()); + } + class PerInstructionStorage { // Abstract class for per-instruction storage. public: @@ -2550,6 +2595,12 @@ class HloPadInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kPad; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), + protobuf_util::ProtobufHashBySerialization(padding_config())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2792,6 +2843,10 @@ class HloIotaInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kIota; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine(std::move(h), iota_dimension()); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; @@ -2861,6 +2916,13 @@ class HloDotInstruction : public HloInstruction { return hlo->opcode() == HloOpcode::kDot; } + void HashAdditionalAttributes(absl::HashState h) const override { + absl::HashState::combine( + std::move(h), + protobuf_util::ProtobufHashBySerialization(dot_dimension_numbers()), + protobuf_util::ProtobufHashBySerialization(precision_config())); + } + private: void PrintExtraAttributesImpl(AttributePrinter& printer, const HloPrintOptions& options) const override; diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index e111be91e2c82f..d0c98dc07d5ba0 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -1414,10 +1414,12 @@ xla_cc_test( deps = [ ":sort_simplifier", "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/hlo/testlib:test", "//xla/service:pattern_matcher", + "@com_google_absl//absl/status:status_matchers", "@com_google_googletest//:gtest_main", "@tsl//tsl/platform:statusor", ], diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc index 436c6643f4d9b1..72d8948f43ba83 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/sort_simplifier_test.cc @@ -17,7 +17,9 @@ limitations under the License. #include +#include "absl/status/status_matchers.h" #include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/test.h" @@ -27,6 +29,7 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::IsOkAndHolds; namespace m = match; using SortSimplifierTest = HloHardwareIndependentTestBase; @@ -50,14 +53,15 @@ TEST_F(SortSimplifierTest, RemoveUnusedSortOperandArrayResult) { dimensions={1}, to_apply=compare ROOT gte = f32[64,8732]{1,0} get-tuple-element(sort), index=0 })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; uint64_t num_executions = 0; + bool changed = false; do { num_executions++; - } while (simplifier.Run(module.get()).value()); + ASSERT_OK_AND_ASSIGN(changed, simplifier.Run(module.get())); + } while (changed); EXPECT_EQ(num_executions, 2); auto root = module->entry_computation()->root_instruction(); EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(0)))); @@ -88,11 +92,10 @@ TEST_F(SortSimplifierTest, RemoveUnusedSortOperandTuple) { gte.1 = u32[64,87] get-tuple-element(sort), index=2 ROOT tuple = (f32[64,87], u32[64,87]) tuple(gte.0, gte.1) })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; - EXPECT_TRUE(simplifier.Run(module.get()).value()); + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(true)); auto root = module->entry_computation()->root_instruction(); EXPECT_THAT( root, @@ -119,11 +122,10 @@ TEST_F(SortSimplifierTest, DontRemoveUnusedSortKey) { sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), dimensions={1}, to_apply=compare ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; - EXPECT_FALSE(simplifier.Run(module.get()).value()); + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(false)); } TEST_F(SortSimplifierTest, RemoveUnusedFirstOperand) { @@ -145,17 +147,83 @@ TEST_F(SortSimplifierTest, RemoveUnusedFirstOperand) { dimensions={1}, to_apply=compare ROOT gte = s32[64,8732]{1,0} get-tuple-element(sort), index=1 })"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); SortSimplifier simplifier; uint64_t num_executions = 0; + bool changed = false; do { num_executions++; - } while (simplifier.Run(module.get()).value()); + ASSERT_OK_AND_ASSIGN(changed, simplifier.Run(module.get())); + } while (changed); EXPECT_EQ(num_executions, 2); auto root = module->entry_computation()->root_instruction(); EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(1)))); } + +TEST_F(SortSimplifierTest, DontRemoveUnusedSortOperandWhenSortIsRoot) { + const char* hlo_string = R"( + HloModule sort_root + + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT + } + + ENTRY sort_computation { + keys = f32[64] parameter(0) + values = s32[64] parameter(1) + ROOT sort = (f32[64], s32[64]) sort(keys, values), + dimensions={0}, to_apply=compare + })"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(false)); +} + +TEST_F(SortSimplifierTest, DoesNotRemoveOldComparatorWhenShared) { + const char* hlo_string = R"( + HloModule shared_comparator + + // CHECK-DAG: %compare ({{.*}}p.0.lhs: f32[], {{.*}}p.0.rhs: f32[], {{.*}}p.1.lhs: s32[], {{.*}}p.1.rhs: s32[]) -> pred[] + // CHECK-DAG: %compare.clone ({{.*}}) -> pred[] + compare { + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT + } + + // CHECK: ENTRY %sort_computation + ENTRY sort_computation { + keys1 = f32[64] parameter(0) + values1 = s32[64] parameter(1) + // CHECK: %[[NEW_SORT:.*]] = f32[64]{{.*}}sort(%keys1), dimensions={0}, to_apply=%compare.clone + sort1 = (f32[64], s32[64]) sort(keys1, values1), + dimensions={0}, to_apply=compare + gte1 = f32[64] get-tuple-element(sort1), index=0 + + keys2 = f32[64] parameter(2) + values2 = s32[64] parameter(3) + // CHECK: %sort2 = (f32[64]{{.*}}, s32[64]{{.*}}) sort(%keys2, %values2), dimensions={0}, to_apply=%compare + sort2 = (f32[64], s32[64]) sort(keys2, values2), + dimensions={0}, to_apply=compare + gte2.0 = f32[64] get-tuple-element(sort2), index=0 + gte2.1 = s32[64] get-tuple-element(sort2), index=1 + + // CHECK: ROOT %tuple = (f32[64]{{.*}}, f32[64]{{.*}}, s32[64]{{.*}}) tuple(%[[NEW_SORT]], %sort2#0, %sort2#1) + ROOT tuple = (f32[64], f32[64], s32[64]) tuple(gte1, gte2.0, gte2.1) + })"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + SortSimplifier simplifier; + EXPECT_THAT(simplifier.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} } // namespace } // namespace xla diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td index d7ce2fdd56487e..28bba3d301c3d1 100644 --- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td +++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_common.td @@ -22,6 +22,7 @@ def MHLO_Dialect : Dialect { let useDefaultAttributePrinterParser = 0; let useDefaultTypePrinterParser = 0; + let useStrictPropertiesInAssemblyFormat = 0; } include "mhlo/IR/hlo_base.td" diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td index 730141fdda7f5d..3833f81c0dfc70 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td @@ -24,6 +24,7 @@ def TPU_Dialect : Dialect { let cppNamespace = "::mlir::tpu"; let useDefaultAttributePrinterParser = 1; let useDefaultTypePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; let extraClassDeclaration = [{ static constexpr StringRef GetMainKey() { return "tpu.main"; } static constexpr StringRef GetCoreTypeKey() { return "tpu.core_type"; } diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc index efcd52db06227a..971c42c5673635 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -1638,6 +1639,42 @@ TEST(StreamExecutorGpuClientTest, CopyFromPinnedHostMemorySpace) { *literal)); } +namespace { + +class FailingHostMemoryAllocator : public HostMemoryAllocator { + public: + OwnedPtr Allocate(size_t size, const AllocateOptions& options) override { + return nullptr; + } +}; + +} // namespace + +TEST(StreamExecutorGpuClientTest, + ToLiteralReturnsResourceExhaustedWhenHostStagingPoolIsFull) { + GpuClientOptions options = GetTestGpuClientOptions(); + options.host_memory_allocator_factory = + [](HostMemoryAllocator::Options options) + -> absl::StatusOr> { + return std::make_unique(); + }; + + ASSERT_OK_AND_ASSIGN(auto client, GetStreamExecutorGpuClient(options)); + ASSERT_OK_AND_ASSIGN( + auto* device_memory_space, + client->addressable_devices()[0]->default_memory_space()); + + constexpr int64_t kTransferBytes = int64_t{1} << 20; + ASSERT_OK_AND_ASSIGN( + auto device_buffer, + client->CreateUninitializedBuffer( + ShapeUtil::MakeShape(U8, {kTransferBytes}), device_memory_space)); + + EXPECT_THAT(device_buffer->ToLiteral().Await(), + StatusIs(absl::StatusCode::kResourceExhausted, + HasSubstr("host staging buffer"))); +} + TEST(StreamExecutorGpuClientTest, CopyToPinnedHostMemorySpaceInt4) { TF_ASSERT_OK_AND_ASSIGN( auto client, GetStreamExecutorGpuClient(GetTestGpuClientOptions())); diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index aa7ae99cc0a108..036dbe82c319f0 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -424,10 +424,6 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBuffer( return absl::InternalError( "se::DeviceAddressAllocator is null in PjRtStreamExecutorRawClient."); } - ABSL_ASSIGN_OR_RETURN( - auto buffer, - allocator_->Allocate(local_device->local_device_id().value(), - on_device_bytes_count, true, layout_memory_space)); tsl::AsyncValueRef mem; if (has_custom_host_memory_allocator_ && layout_memory_space == Layout::kHostMemorySpace) { @@ -437,10 +433,22 @@ absl::StatusOr PjRtStreamExecutorRawClient::AllocateRawBuffer( }; auto buffer = GetHostMemoryAllocator()->Allocate(on_device_bytes_count, alloc_opts); + if (on_device_bytes_count > 0 && buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte buffer from the custom host memory " + "allocator. The allocator may be exhausted or fragmented, or the " + "underlying pinned allocation failed; check earlier allocator " + "warnings for the root cause.", + on_device_bytes_count); + } se::DeviceAddressBase address(buffer.get(), on_device_bytes_count); mem = RawSEDeviceMemory::CreateForeign(address, [buffer = std::move(buffer)]() {}); } else { + ABSL_ASSIGN_OR_RETURN( + auto buffer, + allocator_->Allocate(local_device->local_device_id().value(), + on_device_bytes_count, true, layout_memory_space)); mem = RawSEDeviceMemory::Create(buffer.Release(), local_device, allocator_); } if (client_ != nullptr && local_device->allocation_model() != @@ -582,6 +590,15 @@ PjRtStreamExecutorClient::AllocateLinearizeDest( alloc_opts.local_device_id = local_device->local_device_id(); HostMemoryAllocator::OwnedPtr staging_buffer = GetHostMemoryAllocator()->Allocate(size, alloc_opts); + if (size > 0 && staging_buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for a " + "host-to-device transfer. The pinned host pool may be exhausted or " + "fragmented (see XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the " + "underlying pinned allocation failed; check earlier allocator " + "warnings for the root cause.", + size); + } absl::Span span(staging_buffer.get(), size); return PjRtStagingBuffer::Create( diff --git a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc index 3710029a95e619..eba194a6309659 100644 --- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc +++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc @@ -147,6 +147,16 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent( alloc_opts.local_device_id = local_device->local_device_id(); staging_buffer = client->GetHostMemoryAllocator()->Allocate( transfer_size, alloc_opts); + if (staging_buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for " + "a host-to-device transfer. The pinned host pool may be " + "exhausted or fragmented (see " + "XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the underlying pinned " + "allocation failed; check earlier allocator warnings for the " + "root cause.", + transfer_size); + } auto copy_to_staging_buffer = [src, transfer_size, staging_buffer]() mutable { tsl::profiler::TraceMe trace("H2D Copy To Staging Buffer"); @@ -218,6 +228,16 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent( alloc_opts.local_device_id = local_device->local_device_id(); staging_buffer = client->GetHostMemoryAllocator()->Allocate( transfer_size, alloc_opts); + if (staging_buffer == nullptr) { + return ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for " + "a device-to-host transfer. The pinned host pool may be " + "exhausted or fragmented (see " + "XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the underlying pinned " + "allocation failed; check earlier allocator warnings for the " + "root cause.", + transfer_size); + } ABSL_RETURN_IF_ERROR( stream->Memcpy(staging_buffer.get(), sub_buffer, transfer_size)); auto copy_from_staging_buffer = [dst, transfer_size, @@ -367,6 +387,18 @@ void PjRtStreamExecutorRawBuffer::CopyTo( std::shared_ptr staging_buffer = client_->GetHostMemoryAllocator()->Allocate(GetOnDeviceSizeInBytes(), alloc_opts); + if (GetOnDeviceSizeInBytes() > 0 && staging_buffer == nullptr) { + absl::Status s = ResourceExhausted( + "Failed to allocate a %d-byte pinned host staging buffer for a " + "device-to-device transfer. The pinned host pool may be exhausted or " + "fragmented (see XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB), or the " + "underlying pinned allocation failed; check earlier allocator " + "warnings for the root cause.", + GetOnDeviceSizeInBytes()); + definition_event_promise.SetError(s); + src_usage_event_promise.SetError(s); + return; + } auto d2h_event = CopyRawDeviceToHostAndReturnEvent( staging_buffer.get(), 0, GetOnDeviceSizeInBytes(), {}); if (!d2h_event.ok()) { diff --git a/third_party/xla/xla/python/_profiler.pyi b/third_party/xla/xla/python/_profiler.pyi index 8b2e09aec6e49d..8f40c019618fec 100644 --- a/third_party/xla/xla/python/_profiler.pyi +++ b/third_party/xla/xla/python/_profiler.pyi @@ -33,7 +33,7 @@ class ProfilerSession: def stop_and_export(self, arg: str, /) -> None: ... def stop(self) -> bytes: ... - def stop_and_get_profile_data() -> ProfileData: ... + def stop_and_get_profile_data(self) -> ProfileData: ... def export(self, arg0: bytes, arg1: str, /) -> None: ... class ProfileOptions: @@ -97,6 +97,11 @@ class TraceMe: @staticmethod def is_enabled() -> bool: ... +def export_to_xprof( + profile_data: "tensorflow::profiler::python::ProfileData", + xprof_dir: str, + session_id: str = ..., +) -> None: ... def get_profiled_instructions_proto(tensorboard_dir: str) -> bytes: ... def get_instructions_profile( tensorboard_dir: str, diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td b/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td index f5222425d8166d..d8b5b35ab39715 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_dialect.td @@ -31,6 +31,7 @@ def Ifrt_Dialect : Dialect { let cppNamespace = "::xla::ifrt"; let useDefaultTypePrinterParser = 1; let useDefaultAttributePrinterParser = 1; + let useStrictPropertiesInAssemblyFormat = 0; let hasOperationAttrVerify = 1; let hasRegionArgAttrVerify = 1; } diff --git a/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td b/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td index 299a9f7fe1b8f3..b4a968745a1f75 100644 --- a/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td +++ b/third_party/xla/xla/python/ifrt/ir/vifrt_dialect.td @@ -42,6 +42,7 @@ def Vifrt_Dialect : Dialect { let useDefaultAttributePrinterParser = 0; let useDefaultTypePrinterParser = 0; + let useStrictPropertiesInAssemblyFormat = 0; } //===--------------------------------------------------------------------------- diff --git a/third_party/xla/xla/python/profiler.cc b/third_party/xla/xla/python/profiler.cc index 3bd55afee08a40..d4ba81a74e9943 100644 --- a/third_party/xla/xla/python/profiler.cc +++ b/third_party/xla/xla/python/profiler.cc @@ -256,13 +256,24 @@ NB_MODULE(_profiler, m) { "stop_and_get_profile_data", [](ProfilerSessionWrapper* sess) -> tensorflow::profiler::python::ProfileData { + // Serialize chunked XSpaces when continuous profiling is enabled to + // avoid hitting the 2GB single-protobuf limit, and pass session_id. + if (sess->session->IsContinuousProfilingEnabled()) { + xla::ThrowIfError(sess->session->Stop()); + auto xspaces = + std::make_shared>( + sess->session->SerializeChunks()); + return tensorflow::profiler::python::ProfileData( + std::move(xspaces), sess->session_id); + } auto xspace = std::make_shared(); // Disables the ProfilerSession xla::ThrowIfError(sess->session->CollectData(xspace.get())); - return tensorflow::profiler::python::ProfileData(xspace); + return tensorflow::profiler::python::ProfileData(std::move(xspace), + sess->session_id); }, nb::call_guard(), - nb::sig("def stop_and_get_profile_data() -> ProfileData")) + nb::sig("def stop_and_get_profile_data(self) -> ProfileData")) .def("export", [](ProfilerSessionWrapper* sess, nb::bytes xspace, const std::string& tensorboard_dir) { tensorflow::profiler::XSpace xspace_proto; @@ -366,6 +377,41 @@ NB_MODULE(_profiler, m) { .def("set_metadata", &TraceMeWrapper::SetMetadata) .def_static("is_enabled", &TraceMeWrapper::IsEnabled); + m.def( + "export_to_xprof", + [](const tensorflow::profiler::python::ProfileData& profile_data, + const std::string& xprof_dir, const std::string& session_id = "") { + // Fall back to profile_data.session_id() when session_id is empty, + // and export chunked XSpaces if continuous profiling was used. + const std::string& effective_session_id = + session_id.empty() ? profile_data.session_id() : session_id; + nb::gil_scoped_release release; + if (profile_data.raw_xspaces() != nullptr) { + if (effective_session_id.empty()) { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + xprof_dir, *profile_data.raw_xspaces())); + } else { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + xprof_dir, effective_session_id, *profile_data.raw_xspaces())); + } + return; + } + if (profile_data.raw_xspace() == nullptr) { + throw xla::XlaRuntimeError("ProfileData does not contain XSpace"); + } + if (effective_session_id.empty()) { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + *profile_data.raw_xspace(), xprof_dir, + /* also_export_trace_json= */ true)); + } else { + xla::ThrowIfError(tsl::profiler::ExportToTensorBoard( + *profile_data.raw_xspace(), xprof_dir, effective_session_id, + /* also_export_trace_json= */ true)); + } + }, + nb::arg("profile_data"), nb::arg("xprof_dir"), + nb::arg("session_id") = ""); + m.def( "get_profiled_instructions_proto", [](std::string tensorboard_dir) { diff --git a/third_party/xla/xla/python/profiler/profile_data_lib.cc b/third_party/xla/xla/python/profiler/profile_data_lib.cc index bbb0a12cef3b0f..08aca8b8eebe72 100644 --- a/third_party/xla/xla/python/profiler/profile_data_lib.cc +++ b/third_party/xla/xla/python/profiler/profile_data_lib.cc @@ -211,9 +211,19 @@ ProfileData::ProfileData(const void* serialized_xspace_ptr, } } -/*explicit*/ ProfileData::ProfileData(std::shared_ptr xspace_ptr) { - xspace_ = xspace_ptr; -} +// Store session_id on ProfileData and support chunked XSpaces while +// aliasing the first chunk so plane visitors remain valid without null derefs. +/*explicit*/ ProfileData::ProfileData(std::shared_ptr xspace_ptr, + std::string session_id) + : xspace_(std::move(xspace_ptr)), session_id_(std::move(session_id)) {} + +/*explicit*/ ProfileData::ProfileData( + std::shared_ptr> xspaces_ptr, std::string session_id) + : xspace_((xspaces_ptr && !xspaces_ptr->empty()) + ? std::shared_ptr(xspaces_ptr, &(*xspaces_ptr)[0]) + : std::make_shared()), + xspaces_(std::move(xspaces_ptr)), + session_id_(std::move(session_id)) {} /*explicit*/ ProfileData::ProfileData(const nb::bytes& serialized_xspace) { if (!xspace_) { diff --git a/third_party/xla/xla/python/profiler/profile_data_lib.h b/third_party/xla/xla/python/profiler/profile_data_lib.h index cdfeee940e30e2..9051d4b143c06c 100644 --- a/third_party/xla/xla/python/profiler/profile_data_lib.h +++ b/third_party/xla/xla/python/profiler/profile_data_lib.h @@ -167,7 +167,13 @@ class ProfileData { ProfileData(const void* serialized_xspace_ptr, size_t serialized_xspace_size); - explicit ProfileData(std::shared_ptr xspace_ptr); + // Supports chunked XSpace vectors (continuous profiling >2GB) and + // preserves session_id on ProfileData for standalone export_profile_data(). + explicit ProfileData(std::shared_ptr xspace_ptr, + std::string session_id = ""); + + explicit ProfileData(std::shared_ptr> xspaces_ptr, + std::string session_id = ""); explicit ProfileData(const nb::bytes& serialized_xspace); @@ -177,11 +183,21 @@ class ProfileData { ProfilePlane* find_plane_with_name(const std::string& name) const; + std::shared_ptr raw_xspace() const { return xspace_; } + + // Exposes chunked XSpaces and session_id so export_to_xprof can save >2GB + // continuous profiles and retain session_id when exporting ProfileData. + std::shared_ptr> raw_xspaces() const { return xspaces_; } + + const std::string& session_id() const { return session_id_; } + private: // The actual XSpace protobuf we are wrapping around. A shared ptr is used so // the different levels of visitors (ProfileData, ProfilePlane, // ProfileLine, etc.) don't depend on the lifetime of others. std::shared_ptr xspace_; + std::shared_ptr> xspaces_; + std::string session_id_; }; ProfileData from_serialized_xspace(const std::string& serialized_xspace); diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 281accbe3b6b0e..0d21749f924035 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2903,6 +2903,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], @@ -3541,7 +3542,7 @@ cc_library( deps = [ ":computation_layout", ":hlo_graph_dumper", - ":logical_buffer", + ":hlo_value", ":while_util", "//xla:permutation_util", "//xla:shape_layout", @@ -3549,7 +3550,7 @@ cc_library( "//xla:status_macros", "//xla:util", "//xla:xla_data_proto_cc", - "//xla/hlo/analysis:tuple_points_to_analysis", + "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/hlo/transforms/simplifiers:hlo_dce", @@ -3682,6 +3683,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@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", @@ -3871,9 +3873,9 @@ xla_cc_test( srcs = ["layout_assignment_test.cc"], deps = [ ":computation_layout", + ":hlo_value", ":hlo_verifier", ":layout_assignment", - ":logical_buffer", ":pattern_matcher", "//xla:literal", "//xla:literal_util", diff --git a/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc b/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc index c2b19168678478..7950039e9d08ad 100644 --- a/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc +++ b/third_party/xla/xla/service/async_collective_custom_call_rewriter_test.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_computation.h" @@ -37,6 +38,8 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::StatusIs; + class AsyncCollectiveCustomCallRewriterTest : public HloHardwareIndependentTestBase {}; @@ -414,10 +417,10 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, std::unique_ptr module, ParseAndReturnUnverifiedModule(unsupported_instruction_hlo)); AsyncCollectiveCustomCallRewriter rewriter(use_legacy_collectives); - EXPECT_THAT(rewriter.Run(module.get()), - ::testing::status::StatusIs( - absl::StatusCode::kUnimplemented, - ::testing::HasSubstr("contains a negate instruction"))); + EXPECT_THAT( + rewriter.Run(module.get()), + StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("contains a negate instruction"))); } // Both modes require the all-gather to have exactly one operand. @@ -428,10 +431,9 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, std::unique_ptr module, ParseAndReturnUnverifiedModule(variadic_all_gather_hlo)); AsyncCollectiveCustomCallRewriter rewriter(use_legacy_collectives); - EXPECT_THAT( - rewriter.Run(module.get()), - ::testing::status::StatusIs(absl::StatusCode::kUnimplemented, - ::testing::HasSubstr("has 2 operands"))); + EXPECT_THAT(rewriter.Run(module.get()), + StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("has 2 operands"))); } // Both modes require the compute-on computation to contain exactly one @@ -443,9 +445,8 @@ TEST_F(AsyncCollectiveCustomCallRewriterTest, ParseAndReturnUnverifiedModule(two_all_gathers_hlo)); AsyncCollectiveCustomCallRewriter rewriter(use_legacy_collectives); EXPECT_THAT(rewriter.Run(module.get()), - ::testing::status::StatusIs( - absl::StatusCode::kUnimplemented, - ::testing::HasSubstr("contains at least two"))); + StatusIs(absl::StatusCode::kUnimplemented, + ::testing::HasSubstr("contains at least two"))); } } diff --git a/third_party/xla/xla/service/copy_insertion.cc b/third_party/xla/xla/service/copy_insertion.cc index bb85d255e64bc5..0332767a7e9c9d 100644 --- a/third_party/xla/xla/service/copy_insertion.cc +++ b/third_party/xla/xla/service/copy_insertion.cc @@ -632,19 +632,39 @@ bool IsDisjointInPlaceWhileTupleIndex(const HloInstruction* while_instr, return curr == single_reader || curr == param_gte; } -absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, +// Only such a while makes AddCopiesForWhile read HloBuffers. +// LINT.IfChange(disabled_copies_prescan) +bool HasWhileWithDisabledCopies(const HloModule& module) { + for (const HloComputation* computation : module.computations()) { + for (const HloInstruction* instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::kWhile && + HasDisableWhileLoopCopiesAttr(instruction)) { + return true; + } + } + } + return false; +} +// LINT.ThenChange(:disabled_copies_buffer_reads) + +// `alias_analysis` is only consulted for a while annotated with +// xla_disable_while_loop_copies and may be null otherwise. +absl::Status AddCopiesForWhile(const HloDataflowAnalysis& dataflow, + const HloAliasAnalysis* alias_analysis, HloInstruction* xla_while) { VLOG(2) << "Adding copies for kWhile instruction " << xla_while->name(); TF_RET_CHECK(xla_while->opcode() == HloOpcode::kWhile); ShapeTree indices_to_copy(xla_while->shape()); - IndicesToCopyForWhile(alias_analysis.dataflow_analysis(), xla_while, - &indices_to_copy); + IndicesToCopyForWhile(dataflow, xla_while, &indices_to_copy); + // LINT.IfChange(disabled_copies_buffer_reads) if (HasDisableWhileLoopCopiesAttr(xla_while)) { VLOG(2) << "While loop copy insertion partially disabled via frontend " "attribute for " << xla_while->name(); + TF_RET_CHECK(alias_analysis != nullptr); + // LINT.ThenChange(:disabled_copies_prescan) // Clear copy requirement for scalars, async pipelined, and disjoint // in-place tuple indices. for (auto& pair : indices_to_copy) { @@ -670,7 +690,7 @@ absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, return; } std::vector buffers = - alias_analysis.ComputeBuffersAt(while_init, index); + alias_analysis->ComputeBuffersAt(while_init, index); for (const HloBuffer* buffer : buffers) { if (!seen_buffers.insert(buffer).second) { *init_indices_to_copy.mutable_element(index) = true; @@ -746,9 +766,8 @@ absl::Status AddCopiesForWhile(const HloAliasAnalysis& alias_analysis, // Add copies for the operands of in-place operations. RemoveUnnecessaryCopies // will remove the unnecessary copies. -absl::Status AddCopiesForInPlaceOperation( - const HloAliasAnalysis& alias_analysis, HloInstruction* in_place_op, - int64_t operand_number) { +absl::Status AddCopiesForInPlaceOperation(HloInstruction* in_place_op, + int64_t operand_number) { VLOG(2) << "Adding copies for in-place operation " << in_place_op->name(); HloInstruction* operand = in_place_op->mutable_operand(operand_number); ABSL_ASSIGN_OR_RETURN(HloInstruction * deep_copy, @@ -874,13 +893,12 @@ absl::Status StripControlDependenciesFrom(HloInstruction* instruction) { // roots, in order to resolve interference. We later rely on // RemoveUnnecessaryCopies to drop the unnecessary ones. absl::Status CopyInsertion::AddCopiesForConditional( - const HloAliasAnalysis& alias_analysis, HloInstruction* conditional) { + const HloDataflowAnalysis& dataflow, HloInstruction* conditional) { VLOG(2) << "Adding copies for kConditional instruction " << conditional->name(); ShapeTree indices_to_copy(conditional->shape()); TF_RET_CHECK(conditional->opcode() == HloOpcode::kConditional); - if (!IndicesToCopyForConditional(alias_analysis.dataflow_analysis(), - conditional, &indices_to_copy)) { + if (!IndicesToCopyForConditional(dataflow, conditional, &indices_to_copy)) { VLOG(2) << "No copies necessary for kConditional instruction " << conditional->name(); return absl::OkStatus(); @@ -1236,7 +1254,7 @@ absl::Status AddCopiesForNonCopyableTransitionsRotatedCase( // `chain_start` is the beginning of a non-copyable chain, we add copies to // transition in and out of the chain. absl::Status CopyInsertion::AddCopiesForExplicitNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start) { + HloInstruction* chain_start) { VLOG(2) << "AddCopiesForExplicitNonCopyableTransitions: " << chain_start->ToString(); HloComputation* parent = chain_start->parent(); @@ -1399,12 +1417,11 @@ absl::Status AddCopiesForNonCopyableTransitionsRotatedCase( // (body root) (body root) // absl::Status CopyInsertion::AddCopiesForNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start) { + HloInstruction* chain_start) { if (!IsImplicitNonCopyable(chain_start)) { if (chain_start->IsCustomCall(kPinCustomCallTarget) || chain_start->IsCustomCall(kCreateBufferCustomCallTarget)) { - ABSL_RETURN_IF_ERROR(AddCopiesForExplicitNonCopyableTransitions(alias_analysis, - chain_start)); + ABSL_RETURN_IF_ERROR(AddCopiesForExplicitNonCopyableTransitions(chain_start)); } return absl::OkStatus(); } @@ -1467,8 +1484,27 @@ absl::Status CopyInsertion::AddCopiesForNonCopyableTransitions( absl::Status CopyInsertion::AddCopiesToResolveInterference( HloModule* module, const absl::flat_hash_set& execution_threads) { - ABSL_ASSIGN_OR_RETURN(std::unique_ptr alias_analysis, - HloAliasAnalysis::Run(module, alias_info_)); + // Only the xla_disable_while_loop_copies branch of AddCopiesForWhile reads + // HloBuffers; everything else reads the dataflow analysis. + // HloAliasAnalysis::Run calls HloDataflowAnalysis::Run(*module, + // /*ssa_form=*/true, /*bitcast_defines_value=*/false) and passes alias_info_ + // only to CreateBuffers, so the same two booleans below yield the identical + // dataflow analysis. The analysis is built before the walk on purpose: copies + // inserted by earlier iterations would change the buffers a later build + // would see. + std::unique_ptr alias_analysis; + std::unique_ptr owned_dataflow; + const HloDataflowAnalysis* dataflow = nullptr; + if (HasWhileWithDisabledCopies(*module)) { + ABSL_ASSIGN_OR_RETURN(alias_analysis, + HloAliasAnalysis::Run(module, alias_info_)); + dataflow = &alias_analysis->dataflow_analysis(); + } else { + ABSL_ASSIGN_OR_RETURN(owned_dataflow, + HloDataflowAnalysis::Run(*module, /*ssa_form=*/true, + /*bitcast_defines_value=*/false)); + dataflow = owned_dataflow.get(); + } for (HloComputation* computation : module->MakeNonfusionComputations(execution_threads)) { if (computation->IsAsyncComputation()) { @@ -1477,14 +1513,14 @@ absl::Status CopyInsertion::AddCopiesToResolveInterference( for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { if (instruction->opcode() == HloOpcode::kWhile) { - ABSL_RETURN_IF_ERROR(AddCopiesForWhile(*alias_analysis, instruction)); + ABSL_RETURN_IF_ERROR( + AddCopiesForWhile(*dataflow, alias_analysis.get(), instruction)); } else if (instruction->opcode() == HloOpcode::kConditional) { - ABSL_RETURN_IF_ERROR(AddCopiesForConditional(*alias_analysis, instruction)); + ABSL_RETURN_IF_ERROR(AddCopiesForConditional(*dataflow, instruction)); } else if (IsNonCopyable(instruction)) { // We currently assume that we don't have a custom-call with // output-to-operand aliases for both buffers and non-buffers. - ABSL_RETURN_IF_ERROR( - AddCopiesForNonCopyableTransitions(*alias_analysis, instruction)); + ABSL_RETURN_IF_ERROR(AddCopiesForNonCopyableTransitions(instruction)); } else { // When an operand is a tuple, we avoid copying the operand multiple // times by recording and checking the operand number of operands that @@ -1583,7 +1619,7 @@ absl::Status CopyInsertion::AddCopiesToResolveInterference( } copied_operands.insert(operand_index_in_this_intr); ABSL_RETURN_IF_ERROR(AddCopiesForInPlaceOperation( - *alias_analysis, instruction, operand_index_in_this_intr)); + instruction, operand_index_in_this_intr)); } } } diff --git a/third_party/xla/xla/service/copy_insertion.h b/third_party/xla/xla/service/copy_insertion.h index 31f928a97b84db..16942891bfe432 100644 --- a/third_party/xla/xla/service/copy_insertion.h +++ b/third_party/xla/xla/service/copy_insertion.h @@ -26,6 +26,7 @@ limitations under the License. #include "absl/strings/string_view.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_module.h" @@ -136,15 +137,14 @@ class CopyInsertion : public HloModulePass { // Add copies for conditional instructions. virtual absl::Status AddCopiesForConditional( - const HloAliasAnalysis& alias_analysis, HloInstruction* conditional); + const HloDataflowAnalysis& dataflow, HloInstruction* conditional); // Adds copies for transitioning into and out of non-copyable values. - absl::Status AddCopiesForNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start); + absl::Status AddCopiesForNonCopyableTransitions(HloInstruction* chain_start); // Adds copies for transitioning into and out of non-copyable values for a // explicit non-copyable chain. absl::Status AddCopiesForExplicitNonCopyableTransitions( - const HloAliasAnalysis& alias_analysis, HloInstruction* chain_start); + HloInstruction* chain_start); // Backend specific information about whether an instruction can share buffer // with its operand. @@ -160,6 +160,8 @@ class CopyInsertion : public HloModulePass { const absl::flat_hash_set& execution_threads) override; private: + friend class CopyInsertionTestPeer; + absl::Status AddCopiesToResolveInterference( HloModule* module, const absl::flat_hash_set& execution_threads); diff --git a/third_party/xla/xla/service/copy_insertion_test.cc b/third_party/xla/xla/service/copy_insertion_test.cc index ce2da447ae0a58..0c20d1bfce6910 100644 --- a/third_party/xla/xla/service/copy_insertion_test.cc +++ b/third_party/xla/xla/service/copy_insertion_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "absl/algorithm/container.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" @@ -67,6 +68,17 @@ limitations under the License. namespace op = xla::testing::opcode_matchers; namespace xla { + +// Reaches the private interference step so a test can run it on its own. +class CopyInsertionTestPeer { + public: + static absl::Status AddCopiesToResolveInterference( + CopyInsertion& copy_insertion, HloModule* module) { + return copy_insertion.AddCopiesToResolveInterference( + module, /*execution_threads=*/{}); + } +}; + namespace { using ::testing::NotNull; @@ -1167,6 +1179,125 @@ ENTRY %WhileEntry () -> (f32[16], f32[16]) { EXPECT_EQ(CountCopies(*module), 1); } +// %init_buf reaches the annotated loop twice: directly at index 2 and through +// the output of %first at index 1, which aliases the init of %first. The +// duplicate is only visible on the unmodified module: once %first has its init +// copied, %init_buf no longer shares a buffer with the output of %first. +constexpr absl::string_view kDuplicateInitBuffersThroughEarlierWhileHlo = R"( +HloModule DuplicateInitBuffersThroughEarlierWhile + +%FirstBody (loop_state: (s32[], f32[16])) -> (s32[], f32[16]) { + %loop_state = (s32[], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %c1 = s32[] constant(1) + %next_indvar = s32[] add(%indvar, %c1) + %data = f32[16] get-tuple-element(%loop_state), index=1 + %f1 = f32[] constant(1.0) + %c_add = f32[16] broadcast(%f1) + %next_data = f32[16] add(%data, %c_add) + ROOT %tuple = (s32[], f32[16]) tuple(%next_indvar, %next_data) +} + +%FirstCondition (loop_state: (s32[], f32[16])) -> pred[] { + %loop_state = (s32[], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %limit = s32[] constant(10) + ROOT %cmp = pred[] compare(%indvar, %limit), direction=LT +} + +%Body (loop_state: (s32[], f32[16], f32[16])) -> (s32[], f32[16], f32[16]) { + %loop_state = (s32[], f32[16], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %c1 = s32[] constant(1) + %next_indvar = s32[] add(%indvar, %c1) + %acc1 = f32[16] get-tuple-element(%loop_state), index=1 + %acc2 = f32[16] get-tuple-element(%loop_state), index=2 + %f1 = f32[] constant(1.0) + %update = f32[4] broadcast(%f1) + %c0 = s32[] constant(0) + %dus1 = f32[16] dynamic-update-slice(%acc1, %update, %c0), + frontend_attributes={xla_disjoint_read_write_regions="true"} + %dus2 = f32[16] dynamic-update-slice(%acc2, %update, %c0), + frontend_attributes={xla_disjoint_read_write_regions="true"} + ROOT %tuple = (s32[], f32[16], f32[16]) tuple(%next_indvar, %dus1, %dus2) +} + +%Condition (loop_state: (s32[], f32[16], f32[16])) -> pred[] { + %loop_state = (s32[], f32[16], f32[16]) parameter(0) + %indvar = s32[] get-tuple-element(%loop_state), index=0 + %limit = s32[] constant(10) + ROOT %cmp = pred[] compare(%indvar, %limit), direction=LT +} + +ENTRY %WhileEntry () -> (f32[16], f32[16], f32[16]) { + %c0 = s32[] constant(0) + %indvar_init = s32[] negate(%c0) + %zero = f32[] constant(0.0) + %init_buf = f32[16] broadcast(%zero) + %first_init = (s32[], f32[16]) tuple(%indvar_init, %init_buf) + %first = (s32[], f32[16]) while(%first_init), + condition=%FirstCondition, body=%FirstBody + %first_data = f32[16] get-tuple-element(%first), index=1 + %init_tuple = (s32[], f32[16], f32[16]) tuple(%indvar_init, %first_data, %init_buf) + %while = (s32[], f32[16], f32[16]) while(%init_tuple), + condition=%Condition, body=%Body, + frontend_attributes={xla_disable_while_loop_copies="true"} + %out1 = f32[16] get-tuple-element(%while), index=1 + %out2 = f32[16] get-tuple-element(%while), index=2 + ROOT %root = (f32[16], f32[16], f32[16]) tuple(%out1, %out2, %init_buf) +} +)"; + +TEST_F(WhileCopyInsertionTest, + DisableWhileLoopCopiesDuplicateInitBuffersThroughEarlierWhile) { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule( + kDuplicateInitBuffersThroughEarlierWhileHlo)); + InsertCopies(module.get()); + + EXPECT_EQ(CountCopies(*module->GetComputationWithName("Body")), 0); + // The duplicate at index 2 is copied; %init_buf is live out, so the in-place + // update of index 2 could not run on its buffer. + const HloInstruction* while_hlo = + module->entry_computation()->root_instruction()->operand(0)->operand(0); + EXPECT_THAT(while_hlo->operand(0), + op::Tuple(op::Negate(), op::GetTupleElement(op::While()), + op::Copy(op::Broadcast()))); +} + +// Same module through the interference step alone, so the copy cannot come +// from the alias analysis that later steps build. +TEST_F(WhileCopyInsertionTest, + DisableWhileLoopCopiesDuplicateInitBuffersInterferenceStepOnly) { + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule( + kDuplicateInitBuffersThroughEarlierWhileHlo)); + CopyInsertion copy_insertion(&alias_info_); + ASSERT_OK(CopyInsertionTestPeer::AddCopiesToResolveInterference( + copy_insertion, module.get())); + + EXPECT_EQ(CountCopies(*module->GetComputationWithName("Body")), 0); + // The step deep copies the init, which leaves get-tuple-element/tuple + // rewraps for the tuple simplifier of the full pass; look through them. + auto through_rewraps = [](const HloInstruction* instr) { + while (instr->opcode() == HloOpcode::kGetTupleElement && + instr->operand(0)->opcode() == HloOpcode::kTuple) { + instr = instr->operand(0)->operand(instr->tuple_index()); + } + return instr; + }; + const HloInstruction* while_hlo = + module->entry_computation()->root_instruction()->operand(0)->operand(0); + const HloInstruction* init = while_hlo->operand(0); + ASSERT_EQ(init->opcode(), HloOpcode::kTuple); + EXPECT_THAT(through_rewraps(init->operand(0)), op::Negate()); + EXPECT_THAT(through_rewraps(init->operand(1)), + op::GetTupleElement(op::While())); + const HloInstruction* copy = through_rewraps(init->operand(2)); + ASSERT_THAT(copy, op::Copy()); + EXPECT_THAT(through_rewraps(copy->operand(0)), op::Broadcast()); +} + // Tests Copy Insertion when a while feeds another while // PARAMETER // | | diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 73d9ad7dc320c8..f93d35c2da47d0 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -712,6 +712,7 @@ cc_library( deps = [ ":buffer_allocations", ":gpu_constants", + ":gpu_memory_space_assignment", "//xla:shape_util", "//xla:util", "//xla:xla_proto_cc", @@ -746,6 +747,35 @@ cc_library( ], ) +xla_cc_test( + name = "gpu_executable_buffer_allocator_test", + srcs = ["gpu_executable_buffer_allocator_test.cc"], + deps = [ + ":buffer_allocations", + ":gpu_constants", + ":gpu_executable_buffer_allocator", + ":gpu_memory_space_assignment", + "//xla:executable_run_options", + "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/service:buffer_assignment", + "//xla/service:service_executable_run_options", + "//xla/stream_executor:device_address", + "//xla/stream_executor:device_address_allocator", + "//xla/stream_executor:mock_platform", + "//xla/stream_executor:mock_stream", + "//xla/stream_executor:platform", + "//xla/stream_executor:stream", + "@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/types:span", + "@com_google_googletest//:gtest_main", + ], +) + xla_cc_test( name = "gpu_executable_va_remap_allocator_test", srcs = ["gpu_executable_va_remap_allocator_test.cc"], diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc index bd90c461c433c2..a86b4d8cb936e5 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc +++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc @@ -38,6 +38,7 @@ limitations under the License. #include "xla/service/gpu/buffer_allocations.h" #include "xla/service/gpu/gpu_constants.h" #include "xla/service/gpu/gpu_executable_va_remap_allocator.h" +#include "xla/service/gpu/gpu_memory_space_assignment.h" #include "xla/service/service_executable_run_options.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -52,6 +53,21 @@ limitations under the License. namespace xla::gpu { namespace { +// True for allocations that GenerateBufferAllocations obtains from the device +// memory allocator: everything that is not a parameter, a constant, or a +// thread-local buffer. +bool IsTransientAllocation(const BufferAllocation& allocation) { + return !allocation.is_thread_local() && + !allocation.is_entry_computation_parameter() && + !allocation.is_constant(); +} + +// True for allocations in the collective memory space S(1). +bool IsCollectiveMemoryAllocation(const BufferAllocation& allocation) { + return allocation.color() == + static_cast(MemorySpaceColor::kCollective); +} + absl::Status CheckAlignment(const BufferAllocation& allocation, se::DeviceAddressBase buffer, int arg_idx) { const int64_t expected_alignment = [&] { @@ -230,15 +246,36 @@ GpuExecutableBufferAllocator::ExecutionScope::GenerateBufferAllocations( const int64_t num_buffers = owner_->allocations_.size(); ABSL_RETURN_IF_ERROR(Prepare(run_options, device_ordinal)); - std::vector buffers; - buffers.reserve(num_buffers); + // Resolve allocations in two passes so that every collective memory (S(1)) + // buffer of this execution is allocated before any default memory (S(0)) + // buffer. S(1) buffers back NCCL symmetric windows and can only live in the + // fixed, preallocated part of the shared BFC arena, while S(0) may later be + // allowed to grow past it. Allocating S(1) first guarantees that S(0) + // pressure within the same execution can never take the space S(1) needs. + // Both passes keep allocation-index order, so S(1) placement stays a + // deterministic function of the program, and `buffers` remains indexed by + // allocation index. + std::vector buffers(num_buffers); + auto resolve = [&](int64_t i) -> absl::Status { + const BufferAllocation& allocation = *owner_->allocations_[i]; + ABSL_ASSIGN_OR_RETURN(buffers[i], BufferForAllocation( + get_parameter_buffer, globals, allocation, + memory_allocator, device_ordinal, i)); + return CheckAlignment(allocation, buffers[i], i); + }; + + std::vector deferred_indices; for (int64_t i = 0; i < num_buffers; ++i) { const BufferAllocation& allocation = *owner_->allocations_[i]; - ABSL_ASSIGN_OR_RETURN( - buffers.emplace_back(), - BufferForAllocation(get_parameter_buffer, globals, allocation, - memory_allocator, device_ordinal, i)); - ABSL_RETURN_IF_ERROR(CheckAlignment(allocation, buffers.back(), i)); + if (IsTransientAllocation(allocation) && + !IsCollectiveMemoryAllocation(allocation)) { + deferred_indices.push_back(i); + continue; + } + ABSL_RETURN_IF_ERROR(resolve(i)); + } + for (int64_t i : deferred_indices) { + ABSL_RETURN_IF_ERROR(resolve(i)); } return BufferAllocations(buffers, device_ordinal, memory_allocator); } diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h index e71e1533cc2ea7..a8a8b821e97172 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h +++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h @@ -94,7 +94,9 @@ class GpuExecutableBufferAllocator { // Builds the BufferAllocations for an execution. Entry-computation // parameter buffers are obtained from `get_parameter_buffer`; all other // allocations are resolved internally, including alignment checking and - // any subclass allocation-address policy. + // any subclass allocation-address policy. Collective memory (S(1)) + // buffers are allocated before all other transient buffers, in allocation + // index order within each group. absl::StatusOr GenerateBufferAllocations( const ServiceExecutableRunOptions* run_options, ParameterBufferResolver get_parameter_buffer, diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc new file mode 100644 index 00000000000000..38f9f46121ad86 --- /dev/null +++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator_test.cc @@ -0,0 +1,241 @@ +/* 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/gpu/gpu_executable_buffer_allocator.h" + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/status_matchers.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" +#include "xla/executable_run_options.h" +#include "xla/service/buffer_assignment.h" +#include "xla/service/gpu/buffer_allocations.h" +#include "xla/service/gpu/gpu_constants.h" +#include "xla/service/gpu/gpu_memory_space_assignment.h" +#include "xla/service/service_executable_run_options.h" +#include "xla/shape_util.h" +#include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/device_address_allocator.h" +#include "xla/stream_executor/mock_platform.h" +#include "xla/stream_executor/mock_stream.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/stream.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" + +namespace xla::gpu { +namespace { + +using ::absl_testing::StatusIs; +using ::testing::ElementsAre; +using ::testing::NiceMock; +using ::testing::Pair; + +constexpr int64_t kDefault = static_cast(MemorySpaceColor::kDefault); +constexpr int64_t kCollective = + static_cast(MemorySpaceColor::kCollective); +constexpr int64_t kTempBuffer = + static_cast(MemorySpaceColor::kTempBuffer); + +// Host-backed buffers must satisfy the alignment check on every allocation +// kind. +constexpr uint64_t kAlignment = kXlaAllocatedBufferAlignBytes; + +// Serves every request from aligned host memory and records the memory space +// and size of each request in the order it was made. +class RecordingAllocator final : public se::DeviceAddressAllocator { + public: + RecordingAllocator(const se::Platform* platform, se::Stream* stream) + : se::DeviceAddressAllocator(platform), stream_(stream) {} + + absl::StatusOr> Allocate( + int device_ordinal, uint64_t size, bool retry_on_failure, + int64_t memory_space) override { + if (failing_memory_space_.has_value() && + *failing_memory_space_ == memory_space) { + return absl::ResourceExhaustedError("injected allocation failure"); + } + requests_.push_back({memory_space, size}); + auto storage = std::make_unique(size + kAlignment); + uint8_t* aligned = reinterpret_cast( + (reinterpret_cast(storage.get()) + kAlignment - 1) & + ~uintptr_t{kAlignment - 1}); + storage_.push_back(std::move(storage)); + return se::ScopedDeviceAddress( + se::DeviceAddressBase(aligned, size), device_ordinal, this); + } + + absl::Status Deallocate(int device_ordinal, + se::DeviceAddressBase mem) override { + if (!mem.is_null()) { + ++num_deallocations_; + } + return absl::OkStatus(); + } + + absl::StatusOr GetStream(int device_ordinal) override { + return stream_; + } + + // Makes every subsequent request for `memory_space` fail. + void FailRequestsFor(int64_t memory_space) { + failing_memory_space_ = memory_space; + } + + // (memory_space, size) of every successful request, in request order. + const std::vector>& requests() const { + return requests_; + } + + int num_deallocations() const { return num_deallocations_; } + + private: + se::Stream* stream_; + std::optional failing_memory_space_; + std::vector> requests_; + std::vector> storage_; + int num_deallocations_ = 0; +}; + +class GpuExecutableBufferAllocatorTest : public ::testing::Test { + protected: + GpuExecutableBufferAllocatorTest() : allocator_(&platform_, &stream_) { + run_options_.set_stream(&stream_); + service_run_options_ = ServiceExecutableRunOptions(run_options_); + } + + // Builds the BufferAllocations for one execution of `allocations` with the + // base (ALWAYS_UPDATE) buffer allocator and the recording device allocator. + absl::StatusOr Generate( + absl::Span allocations) { + GpuExecutableBufferAllocator executable_allocator( + "test", allocations, ShapeUtil::MakeShape(F32, {64}), &debug_options_, + /*thunk_executor=*/nullptr); + auto get_parameter_buffer = [&](const BufferAllocation& allocation) + -> absl::StatusOr { + return GpuExecutableBufferAllocator::ParameterBuffer{ + se::DeviceAddressBase(parameter_storage_, sizeof(parameter_storage_)), + allocation.parameter_number()}; + }; + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr scope, + executable_allocator.CreateExecutionScope(&service_run_options_, + &allocator_, + /*device_ordinal=*/0)); + return scope->GenerateBufferAllocations( + &service_run_options_, get_parameter_buffer, &globals_, &allocator_, + /*device_ordinal=*/0); + } + + NiceMock platform_; + NiceMock stream_; + RecordingAllocator allocator_; + ExecutableRunOptions run_options_; + ServiceExecutableRunOptions service_run_options_; + DebugOptions debug_options_; + GpuExecutableBufferAllocator::BufferAllocToDeviceMemoryMap globals_; + alignas(kAlignment) uint8_t parameter_storage_[1024]; + alignas(kAlignment) uint8_t constant_storage_[256]; +}; + +TEST_F(GpuExecutableBufferAllocatorTest, + AllocatesCollectiveMemoryBeforeOtherTransientBuffers) { + BufferAllocation parameter(/*index=*/0, /*size=*/1024, kDefault); + parameter.set_entry_computation_parameter( + /*parameter_number=*/0, /*param_shape_index=*/{}, + /*parameter_aliased_with_output=*/false); + BufferAllocation default_temp(/*index=*/1, /*size=*/1024, kDefault); + BufferAllocation collective_temp(/*index=*/2, /*size=*/2048, kCollective); + BufferAllocation constant(/*index=*/3, /*size=*/256, kDefault); + constant.set_constant(true); + BufferAllocation collective_output(/*index=*/4, /*size=*/512, kCollective); + collective_output.set_maybe_live_out(true); + BufferAllocation default_output(/*index=*/5, /*size=*/768, kDefault); + default_output.set_maybe_live_out(true); + BufferAllocation thread_local_buffer(/*index=*/6, /*size=*/128, kDefault); + thread_local_buffer.set_is_thread_local(true); + BufferAllocation separate_temp(/*index=*/7, /*size=*/64, kTempBuffer); + std::vector allocations = { + ¶meter, &default_temp, + &collective_temp, &constant, + &collective_output, &default_output, + &thread_local_buffer, &separate_temp}; + globals_[3] = se::DeviceAddressBase(constant_storage_, 256); + + ASSERT_OK_AND_ASSIGN(BufferAllocations buffers, Generate(allocations)); + + // Every collective request precedes every other request, and each group + // keeps allocation index order. + EXPECT_THAT(allocator_.requests(), + ElementsAre(Pair(kCollective, 2048), Pair(kCollective, 512), + Pair(kDefault, 1024), Pair(kDefault, 768), + Pair(kTempBuffer, 64))); + + // The result is still indexed by allocation index. + EXPECT_EQ(buffers.GetDeviceAddress(0).opaque(), + static_cast(parameter_storage_)); + EXPECT_EQ(buffers.GetDeviceAddress(1).size(), 1024); + EXPECT_EQ(buffers.GetDeviceAddress(2).size(), 2048); + EXPECT_EQ(buffers.GetDeviceAddress(3).opaque(), + static_cast(constant_storage_)); + EXPECT_EQ(buffers.GetDeviceAddress(4).size(), 512); + EXPECT_EQ(buffers.GetDeviceAddress(5).size(), 768); + EXPECT_TRUE(buffers.GetDeviceAddress(6).is_null()); + EXPECT_EQ(buffers.GetDeviceAddress(7).size(), 64); + + ASSERT_OK(buffers.TearDown(/*live_addresses=*/{}, allocations)); + EXPECT_EQ(allocator_.num_deallocations(), 5); +} + +TEST_F(GpuExecutableBufferAllocatorTest, + KeepsAllocationIndexOrderWithoutCollectiveMemory) { + BufferAllocation first(/*index=*/0, /*size=*/256, kDefault); + BufferAllocation second(/*index=*/1, /*size=*/512, kDefault); + BufferAllocation third(/*index=*/2, /*size=*/1024, kDefault); + std::vector allocations = {&first, &second, &third}; + + ASSERT_OK_AND_ASSIGN(BufferAllocations buffers, Generate(allocations)); + + EXPECT_THAT(allocator_.requests(), + ElementsAre(Pair(kDefault, 256), Pair(kDefault, 512), + Pair(kDefault, 1024))); + ASSERT_OK(buffers.TearDown(/*live_addresses=*/{}, allocations)); +} + +TEST_F(GpuExecutableBufferAllocatorTest, + ReturnsDefaultMemoryFailureAfterCollectiveMemoryIsAllocated) { + allocator_.FailRequestsFor(kDefault); + BufferAllocation default_temp(/*index=*/0, /*size=*/1024, kDefault); + BufferAllocation collective_temp(/*index=*/1, /*size=*/2048, kCollective); + std::vector allocations = {&default_temp, + &collective_temp}; + + EXPECT_THAT(Generate(allocations).status(), + StatusIs(absl::StatusCode::kResourceExhausted)); + // The collective buffer was requested before the default one failed. + EXPECT_THAT(allocator_.requests(), ElementsAre(Pair(kCollective, 2048))); +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD index d221700bdd30fa..b7800f761d509a 100644 --- a/third_party/xla/xla/service/gpu/model/BUILD +++ b/third_party/xla/xla/service/gpu/model/BUILD @@ -447,6 +447,8 @@ cc_library( "//xla/service/gpu:gpu_fusible", "//xla/service/gpu:launch_dimensions", "//xla/stream_executor:device_description", + "//xla/tsl/concurrency:executor", + "//xla/tsl/concurrency:future", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -580,11 +582,13 @@ cc_library( "//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:tiling_from_block_parameters", "//xla/codegen/xtile/codegen:emitter_helpers", "//xla/hlo/ir:hlo", "//xla/hlo/utils:hlo_traversal", + "//xla/runtime:object_pool", "//xla/service:decision", "//xla/service:hlo_cost_analysis", "//xla/service:instruction_fusion", @@ -592,6 +596,8 @@ cc_library( "//xla/service/gpu:hlo_fusion_analysis", "//xla/service/gpu:ir_emission_utils", "//xla/stream_executor:device_description", + "//xla/tsl/concurrency:executor", + "//xla/tsl/concurrency:future", "//xla/tsl/platform:env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", @@ -640,6 +646,7 @@ xla_cc_test( "//xla/service/gpu:ir_emission_utils", "//xla/stream_executor:device_description", "//xla/tests:xla_internal_test_main", + "//xla/tsl/concurrency:future", "//xla/tsl/platform:env", "//xla/tsl/testing:temporary_directory", "@com_google_absl//absl/container:flat_hash_map", diff --git a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc index 8dc83ba48a520c..1eabdb08f877ee 100644 --- a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc +++ b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.cc @@ -39,6 +39,8 @@ limitations under the License. #include "xla/service/gpu/model/gpu_performance_model_base.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" #include "xla/xla_data.pb.h" namespace xla { @@ -49,12 +51,13 @@ CombinedGpuPerformanceModel::CombinedGpuPerformanceModel( HloFusionAnalysisCache& fusion_analysis_cache, mlir::MLIRContext& mlir_context, HloCostAnalysis::ShapeSizeFunction shape_size, bool use_experimental_tiling, - bool enable_same_shape_multi_output_fusion) + bool enable_same_shape_multi_output_fusion, + MlirContextPool* mlir_context_pool) : device_info_(device_info), fusion_analysis_cache_(fusion_analysis_cache), indexing_model_(&device_info, &fusion_analysis_cache, shape_size, &mlir_context, use_experimental_tiling, - enable_same_shape_multi_output_fusion), + enable_same_shape_multi_output_fusion, mlir_context_pool), model_(device_info, fusion_analysis_cache, cache_) {} absl::StatusOr @@ -262,10 +265,19 @@ void CombinedGpuPerformanceModel::Invalidate( cache_.Invalidate(instruction); } -absl::StatusOr -CombinedGpuPerformanceModel::TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor) { - return indexing_model_.TryFindBestTilingForFusion(fusion_adaptor); +tsl::Future +CombinedGpuPerformanceModel::TryFindBestTilingForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, tsl::Executor* executor) { + return indexing_model_.TryFindBestTilingForFusionAsync(fusion_adaptor, + executor); +} + +tsl::Future +CombinedGpuPerformanceModel::TryFindTopKBestTilingsForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, int top_k, + tsl::Executor* executor) { + return indexing_model_.TryFindTopKBestTilingsForFusionAsync(fusion_adaptor, + top_k, executor); } } // namespace gpu diff --git a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h index 94970dfed3a524..80831d7a1b9df2 100644 --- a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h +++ b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model.h @@ -32,6 +32,8 @@ limitations under the License. #include "xla/service/gpu/model/gpu_performance_model_base.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" namespace xla { namespace gpu { @@ -51,7 +53,8 @@ class CombinedGpuPerformanceModel : public GpuPerformanceModelBase { ABSL_ATTRIBUTE_LIFETIME_BOUND, mlir::MLIRContext& mlir_context ABSL_ATTRIBUTE_LIFETIME_BOUND, HloCostAnalysis::ShapeSizeFunction shape_size, - bool use_experimental_tiling, bool enable_same_shape_multi_output_fusion); + bool use_experimental_tiling, bool enable_same_shape_multi_output_fusion, + MlirContextPool* mlir_context_pool = nullptr); // Returns runtime data analysis results of a single instruction. // @@ -81,8 +84,25 @@ class CombinedGpuPerformanceModel : public GpuPerformanceModelBase { // Returns the best tiling for a fusion. // // DOES NOT cache the result. - absl::StatusOr TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor); + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindBestTilingForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + tsl::Executor* executor = nullptr); + + // Returns top_k best tilings for a fusion. + // + // DOES NOT cache the result. + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindTopKBestTilingsForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + int top_k, + tsl::Executor* executor = nullptr); // Invalidates all cache entries related to given instruction. // diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc index 636570437775f0..3438e5ca5045e3 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc @@ -44,9 +44,11 @@ limitations under the License. #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/MathExtras.h" +#include "mlir/IR/MLIRContext.h" #include "xla/backends/gpu/codegen/triton/fusion.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" @@ -73,6 +75,8 @@ limitations under the License. #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" #include "xla/tsl/platform/env.h" #include "xla/util.h" #include "tsl/platform/path.h" @@ -642,6 +646,231 @@ absl::Status DumpAndLogTopKCandidates( return absl::OkStatus(); } +absl::StatusOr> +FilterAndSortCandidates( + std::vector> raw_candidates, int top_k, + const HloFusionAdaptor& fusion_adaptor) { + absl::InlinedVector candidates; + for (auto& res : raw_candidates) { + if (res.has_value()) { + candidates.push_back(std::move(*res)); + } + } + VLOG(1) << absl::StrCat("Found ", candidates.size(), + " valid tiling candidates."); + absl::c_stable_sort( + candidates, [](const TiledRunTimeData& a, const TiledRunTimeData& b) { + return a.runtime_data.exec_time < b.runtime_data.exec_time; + }); + if (candidates.size() > top_k) { + candidates.resize(top_k); + } + ABSL_RETURN_IF_ERROR(DumpAndLogTopKCandidates(fusion_adaptor, candidates)); + return candidates; +} + +// Evaluates a single tiling candidate: clones `base_tiling_space`, assigns +// `padded_tile_sizes` to the clone, tiles the fusion and estimates its run +// time. Returns std::nullopt if the candidate cannot be tiled or if it +// violates Triton constraints. +// +// If `mlir_context_pool` is not null, the clone is rebound to a context +// borrowed from the pool, so that concurrent evaluations do not contend on a +// single mlir::MLIRContext. +absl::StatusOr> EvaluateCandidate( + const HloFusionAdaptor& fusion_adaptor, + const experimental::TilingSpace& base_tiling_space, + absl::Span padded_tile_sizes, + MlirContextPool* mlir_context_pool, + const se::DeviceDescription& device_info, + HloCostAnalysis::ShapeSizeFunction shape_size, + absl::FunctionRef flops_fn) { + // Cloning is faster than calling TilingSpace::Create() for each tiling + // candidate. + std::optional borrowed_context; + mlir::MLIRContext* target_context = nullptr; + if (mlir_context_pool != nullptr) { + ABSL_ASSIGN_OR_RETURN(borrowed_context, mlir_context_pool->GetOrCreate()); + target_context = (*borrowed_context)->get(); + } + std::unique_ptr tiling_space = + base_tiling_space.Clone(target_context); + ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(padded_tile_sizes)); + VLOG(3) << "Trying tile sizes " << absl::StrJoin(padded_tile_sizes, ","); + + const absl::StatusOr tiled_computation = + experimental::TiledHloComputation::Tile(fusion_adaptor, + std::move(tiling_space)); + if (!tiled_computation.ok()) { + // TODO: b/511080616 - GetValidTilings() must return only tilings that can + // be tiled and we should treat all errors here as a failure. + VLOG(3) << "Tiling failed for " << absl::StrJoin(padded_tile_sizes, ",") + << " with error: " << tiled_computation.status().message(); + return std::nullopt; + } + if (const Decision valid = experimental::VerifyTritonConstraints( + *tiled_computation, device_info); + !valid) { + VLOG(3) << "Triton constraints violated for tiling " << valid.Explain(); + return std::nullopt; + } + + ABSL_ASSIGN_OR_RETURN( + std::optional tiled_run_time_data, + EstimateTiledRunTimeDataImpl(fusion_adaptor, *tiled_computation, + device_info, shape_size, flops_fn)); + + if (tiled_run_time_data.has_value()) { + VLOG(2) << "Accepted tile sizes [" + << absl::StrJoin(tiled_run_time_data->block_level_parameters + .output_tile_sizes.front(), + ",") + << "], exec_time " << tiled_run_time_data->runtime_data.exec_time; + } + return tiled_run_time_data; +} + +tsl::Future TryFindTopKBestTilingsWithTilingSpace( + const HloFusionAdaptor& fusion_adaptor, int top_k, tsl::Executor* executor, + std::function flops_fn, + mlir::MLIRContext* mlir_context, MlirContextPool* mlir_context_pool, + const se::DeviceDescription& device_info, + HloCostAnalysis::ShapeSizeFunction shape_size) { + CHECK(executor == &tsl::InlineExecutor::Instance() || + mlir_context_pool != nullptr || mlir_context->isMultithreadingEnabled()) + << "Concurrent tiling evaluation requires either an MlirContextPool or " + "an MLIRContext with multithreading enabled."; + + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr base_tiling_space_unique, + experimental::TilingSpace::Create(fusion_adaptor, mlir_context)); + + ABSL_ASSIGN_OR_RETURN(auto tilings, base_tiling_space_unique->GetValidTilings()); + VLOG(1) << absl::StrCat( + "TryFindTopKBestTilingsForFusionAsync tiling_space evaluating ", + tilings.size(), " tilings."); + + std::shared_ptr base_tiling_space = + std::move(base_tiling_space_unique); + + std::vector>> candidate_futures; + candidate_futures.reserve(tilings.size()); + for (const llvm::SmallVector& tiling : tilings) { + // Assign padded tile size as Triton emitter will require that. + // Symbolic analysis route hides this assumption. + llvm::SmallVector padded_tile_sizes = + xla::xtile::GetPaddedTileSizes(tiling); + + // Early pruning invalid candidates as an optimization. + if (Decision decision = experimental::VerifySubsetOfTritonConstraints( + padded_tile_sizes, *base_tiling_space, device_info); + !decision) { + VLOG(3) << "Pre-filtering candidate violating Triton constraints: " + << absl::StrJoin(padded_tile_sizes, ",") << ": " + << decision.Explain(); + continue; + } + + if (!DoesTilingFitInRegisters(fusion_adaptor, *base_tiling_space, + padded_tile_sizes, device_info)) { + VLOG(3) << "Pre-filtering candidate exceeding register limit: " + << absl::StrJoin(padded_tile_sizes, ","); + continue; + } + + candidate_futures.push_back(tsl::MakeFutureOn( + *executor, + [base_space = base_tiling_space, &fusion_adaptor, + device_info = &device_info, shape_size, flops_fn, mlir_context_pool, + padded_tile_sizes = std::move(padded_tile_sizes)]() { + return EvaluateCandidate(fusion_adaptor, *base_space, + padded_tile_sizes, mlir_context_pool, + *device_info, shape_size, flops_fn); + })); + } + + if (candidate_futures.empty()) { + return TopKTiledRunTimeDataOrError( + absl::InlinedVector{}); + } + + return tsl::JoinFutures(absl::MakeSpan(candidate_futures)) + .Map(*executor, + [&fusion_adaptor, + top_k](std::vector> results) + -> absl::StatusOr { + ABSL_ASSIGN_OR_RETURN(auto candidates, + FilterAndSortCandidates(std::move(results), top_k, + fusion_adaptor)); + return TopKTiledRunTimeDataOrError(std::move(candidates)); + }); +} + +tsl::Future +TryFindTopKBestTilingsWithSymbolicAnalysis( + const HloFusionAdaptor& fusion_adaptor, int top_k, + std::function flops_fn, + mlir::MLIRContext* mlir_context, const se::DeviceDescription& device_info, + HloCostAnalysis::ShapeSizeFunction shape_size) { + SymbolicTileAnalysisOrError analysis_or_error = + SymbolicTileAnalysis::AnalyzeFusion( + fusion_adaptor, mlir_context, + TritonEmitterConstraints::GetBuilder(device_info)); + + if (const auto* fusion_decision = + std::get_if(&analysis_or_error)) { + VLOG(2) << "TryFindTopKBestTilingsForFusionAsync SymbolicTileAnalysis " + "rejected: " + << fusion_decision->Explain(); + return *fusion_decision; + } + + const auto& analysis = std::get(analysis_or_error); + + ABSL_ASSIGN_OR_RETURN(auto tilings, analysis.GetValidTilings()); + VLOG(1) << absl::StrCat( + "TryFindTopKBestTilingsForFusionAsync symbolic analysis evaluating ", + tilings.size(), " tilings."); + + std::vector> raw_candidates; + raw_candidates.reserve(tilings.size()); + for (const auto& tiling : tilings) { + // TODO(b/372454662): This needs to be adjusted if we want to support more + // than one "real root" (i.e. a root without users). + // Currently ComputeTiledComputation() may fail and return an + // Unimplemented error for cases of multi-output fusion that we do not + // support yet. + auto maybe_tiled_hlo_computation = analysis.ComputeTiledComputation(tiling); + if (!maybe_tiled_hlo_computation.ok()) { + if (maybe_tiled_hlo_computation.status().code() == + absl::StatusCode::kUnimplemented && + absl::StrContains(maybe_tiled_hlo_computation.status().message(), + "multi-output fusion")) { + continue; + } + return maybe_tiled_hlo_computation.status(); + } + + ABSL_ASSIGN_OR_RETURN(std::optional tiled_run_time_data, + EstimateTiledRunTimeDataImpl( + fusion_adaptor, *maybe_tiled_hlo_computation, + device_info, shape_size, flops_fn)); + + if (tiled_run_time_data.has_value()) { + ABSL_ASSIGN_OR_RETURN(FlatTiling flat_tiling, + tiling.Flatten(analysis.GetTilingSpecification())); + VLOG(2) << "Accepted tile sizes [" << absl::StrJoin(flat_tiling, ",") + << "], exec_time " << tiled_run_time_data->runtime_data.exec_time; + raw_candidates.push_back(std::move(tiled_run_time_data)); + } + } + + ABSL_ASSIGN_OR_RETURN(auto candidates, + FilterAndSortCandidates(std::move(raw_candidates), top_k, + fusion_adaptor)); + return TopKTiledRunTimeDataOrError(std::move(candidates)); +} + } // namespace int64_t GpuPerformanceModelWithIndexingAnalysis::FlopsPerElement( @@ -827,207 +1056,90 @@ absl::flat_hash_map PrecomputeFlopsMap( } // namespace internal -absl::StatusOr -GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusion( - const HloFusionAdaptor& fusion_adaptor, int top_k) { +tsl::Future +GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, int top_k, + tsl::Executor* executor) { XLA_SCOPED_LOGGING_TIMER( "GpuPerformanceModelWithIndexingAnalysis::" - "TryFindTopKBestTilingsForFusion"); + "TryFindTopKBestTilingsForFusionAsync"); if (!fusion_adaptor.GetRoots().empty() && fusion_adaptor.GetRoots()[0].instruction().parent() != nullptr) { const HloInstruction& root = fusion_adaptor.GetRoots()[0].instruction(); - VLOG(1) << "TryFindTopKBestTilingsForFusion adaptor root " + VLOG(1) << "TryFindTopKBestTilingsForFusionAsync adaptor root " << root.ToString(HloPrintOptions::ShortParsable()) << " parent " << root.parent()->name(); } - absl::InlinedVector candidates; - - if (use_experimental_tiling_) { - using experimental::TiledHloComputation; - using experimental::TilingSpace; - - ABSL_ASSIGN_OR_RETURN(std::unique_ptr base_tiling_space, - TilingSpace::Create(fusion_adaptor, mlir_context_)); - ABSL_ASSIGN_OR_RETURN(auto tilings, base_tiling_space->GetValidTilings()); - VLOG(1) << absl::StrCat( - "TryFindTopKBestTilingsForFusion tiling_space evaluating ", - tilings.size(), " tilings."); - - const absl::flat_hash_map flops_map = - internal::PrecomputeFlopsMap( - fusion_adaptor, - [this](const HloInstruction* hlo) { return FlopsPerElement(hlo); }); + if (executor == nullptr) { + executor = &tsl::InlineExecutor::Instance(); + } - auto flops_fn = [&flops_map](const HloInstruction* hlo) -> int64_t { + if (use_experimental_tiling_) { + auto flops_map = std::make_shared< + const absl::flat_hash_map>( + internal::PrecomputeFlopsMap(fusion_adaptor, + [this](const HloInstruction* hlo) { + return FlopsPerElement(hlo); + })); + + auto flops_fn = [flops_map](const HloInstruction* hlo) -> int64_t { if (hlo->opcode() == HloOpcode::kParameter || hlo->opcode() == HloOpcode::kTuple || hlo->opcode() == HloOpcode::kGetTupleElement) { return 0; } - auto it = flops_map.find(hlo); - CHECK(it != flops_map.end()) + auto it = flops_map->find(hlo); + CHECK(it != flops_map->end()) << "Missing precomputed FLOPs for " << hlo->name(); return it->second; }; - for (const llvm::SmallVector& tiling : tilings) { - // Assign padded tile size as Triton emitter will require that. - // Symbolic analysis route hides this assumption. - llvm::SmallVector padded_tile_sizes = - xla::xtile::GetPaddedTileSizes(tiling); - - // Early pruning invalid candidates as an optimization. - if (Decision decision = experimental::VerifySubsetOfTritonConstraints( - padded_tile_sizes, *base_tiling_space, *device_info_); - !decision) { - VLOG(3) << "Pre-filtering candidate violating Triton constraints: " - << absl::StrJoin(padded_tile_sizes, ",") << ": " - << decision.Explain(); - continue; - } - - if (!DoesTilingFitInRegisters(fusion_adaptor, *base_tiling_space, - padded_tile_sizes, *device_info_)) { - VLOG(3) << "Pre-filtering candidate exceeding register limit: " - << absl::StrJoin(padded_tile_sizes, ","); - continue; - } - - // Cloning is faster than calling TilingSpace::Create() for each - // tiling candidate. - std::unique_ptr tiling_space = base_tiling_space->Clone(); - ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(padded_tile_sizes)); - VLOG(3) << "Trying tile sizes " << absl::StrJoin(padded_tile_sizes, ","); - - const absl::StatusOr tiled_computation = - TiledHloComputation::Tile(fusion_adaptor, std::move(tiling_space)); - if (!tiled_computation.ok()) { - // TODO: b/511080616 - GetValidTilings() must return only tilings that - // can be tiled and we should treat all errors here as a failure. - VLOG(3) << "Tiling failed for " << absl::StrJoin(tiling, ",") - << " with error: " << tiled_computation.status().message(); - continue; - } - if (const Decision valid = experimental::VerifyTritonConstraints( - *tiled_computation, *device_info_); - !valid) { - VLOG(3) << "Triton constraints violated for tiling " << valid.Explain(); - continue; - } - - ABSL_ASSIGN_OR_RETURN( - std::optional tiled_run_time_data, - EstimateTiledRunTimeDataImpl(fusion_adaptor, *tiled_computation, - *device_info_, shape_size_, flops_fn)); - - if (tiled_run_time_data.has_value()) { - VLOG(2) << "Accepted tile sizes [" - << absl::StrJoin(tiled_run_time_data->block_level_parameters - .output_tile_sizes.front(), - ",") - << "], exec_time " - << tiled_run_time_data->runtime_data.exec_time; - candidates.push_back(*tiled_run_time_data); - } - } - } else { - SymbolicTileAnalysisOrError analysis_or_error = - SymbolicTileAnalysis::AnalyzeFusion( - fusion_adaptor, mlir_context_, - TritonEmitterConstraints::GetBuilder(*device_info_)); - - if (const auto* fusion_decision = - std::get_if(&analysis_or_error)) { - VLOG(2) - << "TryFindTopKBestTilingsForFusion SymbolicTileAnalysis rejected: " - << fusion_decision->Explain(); - return *fusion_decision; - } - - SymbolicTileAnalysis analysis = - std::get(std::move(analysis_or_error)); - - ABSL_ASSIGN_OR_RETURN(auto tilings, analysis.GetValidTilings()); - VLOG(1) << absl::StrCat( - "TryFindTopKBestTilingsForFusion symbolic analysis evaluating ", - tilings.size(), " tilings."); - for (const auto& tiling : tilings) { - // TODO(b/372454662): This needs to be adjusted if we want to support more - // than one "real root" (i.e. a root without users). - // Currently ComputeTiledComputation() may fail and return an - // Unimplemented error for cases of multi-output fusion that we do not - // support yet. - auto maybe_tiled_hlo_computation = - analysis.ComputeTiledComputation(tiling); - if (!maybe_tiled_hlo_computation.ok()) { - if (maybe_tiled_hlo_computation.status().code() == - absl::StatusCode::kUnimplemented && - absl::StrContains(maybe_tiled_hlo_computation.status().message(), - "multi-output fusion")) { - continue; - } - return maybe_tiled_hlo_computation.status(); - } - - ABSL_ASSIGN_OR_RETURN( - std::optional tiled_run_time_data, - EstimateTiledRunTimeDataImpl( - fusion_adaptor, *maybe_tiled_hlo_computation, *device_info_, - shape_size_, [this](const HloInstruction* hlo) { - return FlopsPerElement(hlo); - })); - - if (tiled_run_time_data.has_value()) { - ABSL_ASSIGN_OR_RETURN(FlatTiling flat_tiling, - tiling.Flatten(analysis.GetTilingSpecification())); - VLOG(2) << "Accepted tile sizes [" << absl::StrJoin(flat_tiling, ",") - << "], exec_time " - << tiled_run_time_data->runtime_data.exec_time; - candidates.push_back(*tiled_run_time_data); - } - } + return TryFindTopKBestTilingsWithTilingSpace( + fusion_adaptor, top_k, executor, flops_fn, mlir_context_, + mlir_context_pool_, *device_info_, shape_size_); } - VLOG(1) << absl::StrCat("Found ", candidates.size(), - " valid tiling candidates."); - absl::c_stable_sort( - candidates, [](const TiledRunTimeData& a, const TiledRunTimeData& b) { - return a.runtime_data.exec_time < b.runtime_data.exec_time; - }); - - if (candidates.size() > top_k) { - candidates.resize(top_k); - } - - ABSL_RETURN_IF_ERROR(DumpAndLogTopKCandidates(fusion_adaptor, candidates)); - - return candidates; + auto flops_fn = [this](const HloInstruction* hlo) -> int64_t { + return FlopsPerElement(hlo); + }; + return TryFindTopKBestTilingsWithSymbolicAnalysis(fusion_adaptor, top_k, + flops_fn, mlir_context_, + *device_info_, shape_size_); } -absl::StatusOr -GpuPerformanceModelWithIndexingAnalysis::TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor) { - ABSL_ASSIGN_OR_RETURN(auto top_k_result, TryFindTopKBestTilingsForFusion( - fusion_adaptor, /*top_k=*/1)); - if (std::holds_alternative(top_k_result)) { - return std::get(top_k_result); - } - auto& tilings = - std::get>(top_k_result); - if (tilings.empty()) { - return FusionDecision::Forbid("No valid tilings found."); +tsl::Future +GpuPerformanceModelWithIndexingAnalysis::TryFindBestTilingForFusionAsync( + const HloFusionAdaptor& fusion_adaptor, tsl::Executor* executor) { + if (executor == nullptr) { + executor = &tsl::InlineExecutor::Instance(); } - VLOG(1) - << "TryFindBestTilingForFusion " - << fusion_adaptor.GetRoots().front().instruction().ToString( - HloPrintOptions::ShortParsable()) - << " best output_tile_sizes [" - << absl::StrJoin( - tilings.front().block_level_parameters.output_tile_sizes.front(), - ",") - << "], exec_time " << tilings.front().runtime_data.exec_time; - return tilings.front(); + return TryFindTopKBestTilingsForFusionAsync(fusion_adaptor, /*top_k=*/1, + executor) + .Map(*executor, + [&fusion_adaptor](const TopKTiledRunTimeDataOrError& top_k_result) + -> TiledRunTimeDataOrError { + if (std::holds_alternative(top_k_result)) { + return std::get(top_k_result); + } + const auto& tilings = + std::get>( + top_k_result); + if (tilings.empty()) { + return FusionDecision::Forbid("No valid tilings found."); + } + VLOG(1) + << "TryFindBestTilingForFusionAsync " + << fusion_adaptor.GetRoots().front().instruction().ToString( + HloPrintOptions::ShortParsable()) + << " best output_tile_sizes [" + << absl::StrJoin( + tilings.front() + .block_level_parameters.output_tile_sizes.front(), + ",") + << "], exec_time " << tilings.front().runtime_data.exec_time; + return tilings.front(); + }); } } // namespace gpu diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h index 50cc2989f3cf83..be9d1e10d2e26e 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h @@ -17,6 +17,8 @@ limitations under the License. #define XLA_SERVICE_GPU_MODEL_GPU_INDEXING_PERFORMANCE_MODEL_H_ #include +#include +#include #include #include @@ -31,6 +33,7 @@ limitations under the License. #include "xla/codegen/xtile/block_level_parameters.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/utils/hlo_traversal.h" +#include "xla/runtime/object_pool.h" #include "xla/service/gpu/model/fusion_analysis_cache.h" #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h" #include "xla/service/gpu/model/gpu_performance_model_base.h" @@ -38,10 +41,14 @@ limitations under the License. #include "xla/service/hlo_cost_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/executor.h" +#include "xla/tsl/concurrency/future.h" namespace xla { namespace gpu { +using MlirContextPool = ObjectPool>; + // Contains informations about block level parameters and run time of a fusion. struct TiledRunTimeData { EstimateRunTimeData runtime_data; @@ -72,7 +79,8 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { HloFusionAnalysisCache* fusion_analysis_cache, HloCostAnalysis::ShapeSizeFunction shape_size, mlir::MLIRContext* mlir_context, bool use_experimental_tiling, - bool enable_same_shape_multi_output_fusion) + bool enable_same_shape_multi_output_fusion, + MlirContextPool* mlir_context_pool = nullptr) : hlo_op_profile_(&HloOpProfiles::Singleton().GetProfile(*device_info)), device_info_(device_info), fusion_analysis_cache_(fusion_analysis_cache), @@ -84,6 +92,7 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { /*count_multiple_input_accesses=*/true}, *device_info_), mlir_context_(mlir_context), + mlir_context_pool_(mlir_context_pool), use_experimental_tiling_(use_experimental_tiling), enable_same_shape_multi_output_fusion_( enable_same_shape_multi_output_fusion) {} @@ -125,13 +134,26 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { // Returns FusionDecision if the fusion can't be tiled or there are no valid // block level parameters. // Otherwise returns block level parameters that give the best execution time. - absl::StatusOr TryFindBestTilingForFusion( - const HloFusionAdaptor& fusion_adaptor); + // + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindBestTilingForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + tsl::Executor* executor = nullptr); // Returns top_k (possibly fewer if not enough valid tilings are found) block // level parameters for the given fusion. - absl::StatusOr TryFindTopKBestTilingsForFusion( - const HloFusionAdaptor& fusion_adaptor, int top_k); + // + // NOTE: `fusion_adaptor` and `this` are captured; callers must ensure both + // `fusion_adaptor` and this model instance outlive the returned + // `tsl::Future`. If executor is nullptr, defaults to + // tsl::InlineExecutor::Instance(). + [[nodiscard]] tsl::Future + TryFindTopKBestTilingsForFusionAsync(const HloFusionAdaptor& fusion_adaptor, + int top_k, + tsl::Executor* executor = nullptr); // Returns an estimate how many FLOPs will be used to produce one element of // the output. @@ -144,6 +166,7 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { HloCostAnalysis::ShapeSizeFunction shape_size_; GpuHloCostAnalysis cost_analysis_; mlir::MLIRContext* mlir_context_; + MlirContextPool* mlir_context_pool_; bool use_experimental_tiling_; bool enable_same_shape_multi_output_fusion_; }; 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 711a6ad045107c..03554a5406806c 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 @@ -53,7 +53,9 @@ limitations under the License. #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/concurrency/future.h" #include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/threadpool.h" #include "xla/tsl/testing/temporary_directory.h" #include "xla/util.h" #include "tsl/platform/path.h" @@ -75,7 +77,14 @@ class GpuIndexingPerformanceModelTest bool use_experimental_tiling() const { return GetParam(); } + static std::unique_ptr CreateMlirContext() { + auto ctx = std::make_unique(); + ctx->disableMultithreading(); + return ctx; + } + mlir::MLIRContext mlir_context_; + MlirContextPool mlir_context_pool_{CreateMlirContext, 4}; // The reference times in the test cases below are measured // on A6000 by profiling the execution of the HLOs. se::DeviceDescription device_info_{TestGpuDeviceInfo::RTXA6000DeviceInfo()}; @@ -86,7 +95,8 @@ class GpuIndexingPerformanceModelTest HloCostAnalysis::DefaultShapeSize, &mlir_context_, use_experimental_tiling(), - /*enable_same_shape_multi_output_fusion=*/false}; + /*enable_same_shape_multi_output_fusion=*/false, + &mlir_context_pool_}; size_t WarpSize() const { return ::xla::gpu::WarpSize(device_info_); } @@ -442,7 +452,8 @@ ENTRY entry_computation { ASSERT_OK_AND_ASSIGN( TiledRunTimeDataOrError tiling_result, - indexing_cost_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model_.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); ASSERT_TRUE(std::holds_alternative(tiling_result)); @@ -495,7 +506,8 @@ ENTRY entry_computation { ASSERT_OK_AND_ASSIGN( TiledRunTimeDataOrError tiling_result, - indexing_cost_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model_.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); ASSERT_TRUE(std::holds_alternative(tiling_result)); @@ -538,8 +550,10 @@ ENTRY main { ASSERT_OK_AND_ASSIGN( auto top_k_result, - indexing_cost_model_.TryFindTopKBestTilingsForFusion(*fusion_adaptor, - /*top_k=*/3)); + indexing_cost_model_ + .TryFindTopKBestTilingsForFusionAsync(*fusion_adaptor, + /*top_k=*/3) + .Await()); ASSERT_TRUE((std::holds_alternative>( top_k_result))); @@ -583,7 +597,8 @@ ENTRY main { ASSERT_OK_AND_ASSIGN( TiledRunTimeDataOrError tiling_result, - indexing_cost_model_.TryFindBestTilingForFusion(*fusion_adaptor)); + indexing_cost_model_.TryFindBestTilingForFusionAsync(*fusion_adaptor) + .Await()); ASSERT_TRUE(std::holds_alternative(tiling_result)); @@ -1330,6 +1345,61 @@ ENTRY entry { EXPECT_EQ(flops_map.at(exp), indexing_cost_model_.FlopsPerElement(exp)); } +TEST_P(GpuIndexingPerformanceModelTest, + TryFindTopKBestTilingsForFusionAsync_ParallelMatchesSequential) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +fused_computation { + param_0.1 = f32[1024,1024] parameter(0) + ROOT negate.1 = f32[1024,1024] negate(param_0.1) +} + +ENTRY main { + param_0.3 = f32[1024,1024] parameter(0) + ROOT fusion = f32[1024,1024] fusion(param_0.3), kind=kCustom, calls=fused_computation +} +)")); + + auto fusion_adaptor = HloFusionAdaptor::ForInstruction( + module->entry_computation()->root_instruction()); + + constexpr int top_k = 5; + ASSERT_OK_AND_ASSIGN( + auto seq_result, + indexing_cost_model_ + .TryFindTopKBestTilingsForFusionAsync(*fusion_adaptor, top_k) + .Await()); + + tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), "test", 8); + ASSERT_OK_AND_ASSIGN(auto parallel_result, + indexing_cost_model_ + .TryFindTopKBestTilingsForFusionAsync( + *fusion_adaptor, top_k, thread_pool.AsExecutor()) + .Await()); + + ASSERT_TRUE((std::holds_alternative>( + seq_result))); + const auto& seq_candidates = + std::get>(seq_result); + ASSERT_FALSE(seq_candidates.empty()); + + ASSERT_TRUE((std::holds_alternative>( + parallel_result))); + const auto& parallel_candidates = + std::get>(parallel_result); + + ASSERT_EQ(parallel_candidates.size(), seq_candidates.size()); + for (int i = 0; i < seq_candidates.size(); ++i) { + EXPECT_EQ(parallel_candidates[i].runtime_data.exec_time, + seq_candidates[i].runtime_data.exec_time); + EXPECT_EQ(parallel_candidates[i].block_level_parameters.output_tile_sizes, + seq_candidates[i].block_level_parameters.output_tile_sizes); + EXPECT_EQ(parallel_candidates[i].block_level_parameters.num_warps, + seq_candidates[i].block_level_parameters.num_warps); + } +} + } // namespace } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index 88fa67ba987317..3238f50e4a377e 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -19,12 +19,14 @@ limitations under the License. #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" @@ -36,7 +38,7 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" -#include "xla/hlo/analysis/tuple_points_to_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.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" @@ -51,7 +53,7 @@ limitations under the License. #include "xla/map_util.h" #include "xla/permutation_util.h" #include "xla/service/computation_layout.h" -#include "xla/service/logical_buffer.h" +#include "xla/service/hlo_value.h" #include "xla/service/while_util.h" #include "xla/shape.h" #include "xla/shape_layout.h" @@ -78,7 +80,7 @@ std::ostream& operator<<(std::ostream& out, } BufferLayoutConstraint::BufferLayoutConstraint(const Layout& layout, - const LogicalBuffer& buffer, + const HloValue& buffer, bool mandatory, bool dfs, int64_t priority) : LayoutConstraint(mandatory, dfs, priority), buffer_(&buffer) { @@ -220,75 +222,72 @@ std::string ComputationLayoutConstraint::ToString() const { layout_state_, computation_layout_.ToString()); } -void LayoutAssignment::SetPointsToAnalysis( - std::unique_ptr analysis) { - points_to_analysis_ = std::move(analysis); +void LayoutAssignment::SetDataflowAnalysis( + std::unique_ptr analysis) { + dataflow_analysis_ = std::move(analysis); buffer_sets_cache_.clear(); - buffer_sets_cache_analysis_ = points_to_analysis_.get(); + buffer_sets_cache_analysis_ = dataflow_analysis_.get(); } -PointsToSet::BufferSet* LayoutAssignment::GetBufferSet( +LayoutAssignment::BufferSet* LayoutAssignment::GetBufferSet( const HloInstruction* instruction) const { - DCHECK_EQ(buffer_sets_cache_analysis_, points_to_analysis_.get()) - << "points_to_analysis_ changed behind buffer_sets_cache_; use " - "SetPointsToAnalysis"; + DCHECK_EQ(buffer_sets_cache_analysis_, dataflow_analysis_.get()) + << "dataflow_analysis_ changed behind buffer_sets_cache_; use " + "SetDataflowAnalysis"; auto it = buffer_sets_cache_.find(instruction); if (it != buffer_sets_cache_.end()) { return it->second.get(); } auto& buffer_set = - buffer_sets_cache_ - .emplace(instruction, std::make_unique()) + buffer_sets_cache_.emplace(instruction, std::make_unique()) .first->second; - const auto& points_to_set = points_to_analysis_->GetPointsToSet(instruction); - points_to_set.ForEachElement( - [&buffer_set](const ShapeIndex& /*index*/, - const PointsToSet::BufferList& buffers) { - buffer_set->insert(buffers.begin(), buffers.end()); + const auto& instruction_value_set = + dataflow_analysis_->GetInstructionValueSet(instruction); + instruction_value_set.ForEachElement( + [&buffer_set](const ShapeIndex& /*index*/, const HloValueSet& value_set) { + buffer_set->insert(value_set.values().begin(), + value_set.values().end()); }); return buffer_set.get(); } bool LayoutAssignment::AnyOperandBufferForwarded( const HloInstruction* instruction, int64_t operand_no) const { - // The operand is potentially forwarded if the intersection of points-to sets + // The operand is potentially forwarded if the intersection of value sets // of the operand and the instruction is non-empty. - PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); - PointsToSet::BufferSet* operand_buffers = - GetBufferSet(instruction->operand(operand_no)); + BufferSet* output_buffers = GetBufferSet(instruction); + BufferSet* operand_buffers = GetBufferSet(instruction->operand(operand_no)); // Probe the larger set with elements of the smaller one; intersection // emptiness is symmetric. - PointsToSet::BufferSet* small = output_buffers; - PointsToSet::BufferSet* large = operand_buffers; + BufferSet* small = output_buffers; + BufferSet* large = operand_buffers; if (small->size() > large->size()) { std::swap(small, large); } - return absl::c_any_of( - *small, [&](const LogicalBuffer* b) { return large->count(b) > 0; }); + return absl::c_any_of(*small, + [&](const HloValue* b) { return large->count(b) > 0; }); } bool LayoutAssignment::AllOperandBuffersForwarded( const HloInstruction* instruction, int64_t operand_no) const { - // The operand is potentially forwarded if the intersection of points-to sets + // The operand is potentially forwarded if the intersection of value sets // of the operand and the instruction is non-empty. - PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); - PointsToSet::BufferSet* operand_buffers = - GetBufferSet(instruction->operand(operand_no)); + BufferSet* output_buffers = GetBufferSet(instruction); + BufferSet* operand_buffers = GetBufferSet(instruction->operand(operand_no)); // Each buffer in operand_buffers should also occur in output_buffers. - return absl::c_all_of(*operand_buffers, [&](const LogicalBuffer* b) { + return absl::c_all_of(*operand_buffers, [&](const HloValue* b) { return output_buffers->count(b) > 0; }); } absl::Status LayoutAssignment::SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer, + const HloValue& buffer, bool mandatory, bool dfs, int64_t priority, const HloInstruction* user) { VLOG(3) << "SetBufferLayout : " << buffer << " : " << LayoutUtil::HumanString(layout) << " with priority " << priority << "; mandatory = " << mandatory << "; dfs = " << dfs << "\n"; - ABSL_RETURN_IF_ERROR(points_to_analysis_->VerifyBuffer(buffer)); if (unconstrained_buffer_ids_.erase(buffer.id()) > 0) { VLOG(3) << "Erase buffer from unconstrained ids\n"; } @@ -466,8 +465,8 @@ absl::Status LayoutAssignment::SetInstructionLayout( instruction->shape(), [this, layout, instruction, mandatory, allow_alias, priority]( const Shape& subshape, const ShapeIndex& index) -> absl::Status { - auto buffers = - points_to_analysis_->GetPointsToSet(instruction).element(index); + const auto& buffers = + dataflow_analysis_->GetValueSet(instruction, index).values(); CHECK_EQ(1, buffers.size()); if (!allow_alias) { CHECK_EQ(buffers[0]->instruction(), instruction); @@ -506,8 +505,8 @@ absl::Status LayoutAssignment::SetInstructionLayout( if (!subshape_index.empty() && index != subshape_index) { return absl::OkStatus(); } - auto buffers = - points_to_analysis_->GetPointsToSet(instruction).element(index); + const auto& buffers = + dataflow_analysis_->GetValueSet(instruction, index).values(); CHECK_EQ(1, buffers.size()); if (!allow_alias) { CHECK_EQ(buffers[0]->instruction(), instruction); @@ -538,7 +537,7 @@ absl::Status LayoutAssignment::SetInstructionLayout( } const BufferLayoutConstraint* LayoutAssignment::GetBufferLayoutConstraint( - const LogicalBuffer& buffer) const { + const HloValue& buffer) const { return buffer_constraints_.Find(buffer.id()); } @@ -586,13 +585,18 @@ std::string LayoutAssignment::ToString( "): ", constraints.OperandLayout(instruction, i)->ToString(), "\n"); } } - for (const LogicalBuffer* buffer : - points_to_analysis_->GetBuffersDefinedByInstruction(instruction)) { - auto* buffer_constraint = GetBufferLayoutConstraint(*buffer); - if (buffer_constraint != nullptr) { - absl::StrAppend(&output, " ", buffer->ToString(), " : ", - LayoutUtil::HumanString(buffer_constraint->layout()), - "\n"); + for (const auto& [index, value_set] : + dataflow_analysis_->GetInstructionValueSet(instruction)) { + for (const HloValue* buffer : value_set.values()) { + if (buffer->instruction() != instruction) { + continue; + } + auto* buffer_constraint = GetBufferLayoutConstraint(*buffer); + if (buffer_constraint != nullptr) { + absl::StrAppend(&output, " ", buffer->ToString(), " : ", + LayoutUtil::HumanString(buffer_constraint->layout()), + "\n"); + } } } } @@ -740,7 +744,13 @@ absl::Status LayoutAssignment::AddParameterConstraints( absl::Status LayoutAssignment::AddCollectiveConstraints( HloInstruction* instruction) { - ABSL_RETURN_IF_ERROR(SetInstructionLayout(instruction->shape(), instruction)); + // HloDataflowAnalysis forwards operand HloValues into tuple element {0} of + // async collective start instructions and forwards start HloValues into + // collective done instructions, so allow_alias is required when constraining + // the instruction output shape. + ABSL_RETURN_IF_ERROR(SetInstructionLayout(instruction->shape(), instruction, + /*mandatory=*/true, /*dfs=*/true, + /*allow_alias=*/true)); for (int64_t i = 0; i < instruction->operand_count(); ++i) { CHECK(instruction->shape().IsArray() || (instruction->shape().IsTuple() && @@ -2017,7 +2027,13 @@ void LayoutAssignment::SetupCopiedInstruction(const HloInstruction& instruction, absl::Status LayoutAssignment::CheckLayouts( HloModule* module, const absl::flat_hash_set& execution_threads) { - ABSL_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); + ABSL_ASSIGN_OR_RETURN(auto dataflow_analysis, + HloDataflowAnalysis::Run( + *module, /*ssa_form=*/false, + /*bitcast_defines_value=*/false, execution_threads, + /*propagate_through_calls=*/false, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); for (auto* computation : module->MakeNonfusionComputations(execution_threads)) { for (auto* instruction : computation->instructions()) { @@ -2026,19 +2042,18 @@ absl::Status LayoutAssignment::CheckLayouts( TF_RET_CHECK(LayoutUtil::HasLayout(instruction->shape())); ABSL_RETURN_IF_ERROR(ShapeUtil::ValidateShape(instruction->shape())); - // Use points-to analysis to verify that every subshape element in the - // output of the instruction matches the layout of the logical buffer + // Use dataflow analysis to verify that every subshape element in the + // output of the instruction matches the layout of the HloValue // which could be the source of the subshape value. - const PointsToSet& points_to_set = - points_to_analysis->GetPointsToSet(instruction); - ABSL_RETURN_IF_ERROR(points_to_set.ForEachElementWithStatus( - [&instruction]( - ShapeIndex index, - const PointsToSet::BufferList& buffers) -> absl::Status { + const InstructionValueSet& value_set_tree = + dataflow_analysis->GetInstructionValueSet(instruction); + ABSL_RETURN_IF_ERROR(value_set_tree.ForEachElementWithStatus( + [&instruction](ShapeIndex index, + const HloValueSet& value_set) -> absl::Status { if (ShapeUtil::IsLeafIndex(instruction->shape(), index)) { const Shape& instruction_subshape = ShapeUtil::GetSubshape(instruction->shape(), index); - for (const LogicalBuffer* buffer : buffers) { + for (const HloValue* buffer : value_set.values()) { if (!Shape::Equal() .IgnoreDynamicDimension() .IgnoreBuffer() @@ -2046,7 +2061,7 @@ absl::Status LayoutAssignment::CheckLayouts( buffer->shape())) { return Internal( "Layout of instruction %s at index {%s} does not match " - "source LogicalBuffer %s: %s vs %s", + "source HloValue %s: %s vs %s", instruction->name(), absl::StrJoin(index, ","), buffer->ToString(), ShapeUtil::HumanStringWithLayout(instruction_subshape), @@ -2406,22 +2421,21 @@ absl::Status LayoutAssignment::PropagateConstraints( namespace { // Returns a vector containing all array-shaped uses (instruction and operand -// number) of the given logical buffer or its aliases. +// number) of the given HloValue or its aliases. std::vector> GetArrayUsesOfBuffer( - const TuplePointsToAnalysis::BufferAliasVector& aliases) { + const HloValue& buffer) { std::vector> uses; - for (const auto& buffer_alias : aliases) { - if (!buffer_alias.instruction()->shape().IsArray()) { + for (const HloPosition& position : buffer.positions()) { + if (!position.instruction->shape().IsArray()) { continue; } // This alias must be the top-level (index == {}) of the instruction's // result because the instruction produces an array. - CHECK(buffer_alias.index().empty()); + CHECK(position.index.empty()); // Add all uses of the instruction's output. - for (const HloInstruction* user : buffer_alias.instruction()->users()) { - for (int64_t operand_no : - user->OperandIndices(buffer_alias.instruction())) { + for (const HloInstruction* user : position.instruction->users()) { + for (int64_t operand_no : user->OperandIndices(position.instruction)) { uses.emplace_back(user, operand_no); } } @@ -2437,17 +2451,17 @@ absl::Status LayoutAssignment::PropagateUseConstraintToDefs( const HloInstruction* user) { // Try to set all logical buffers which may be sources of the given operand to // match the given layout. - const PointsToSet& points_to_set = - points_to_analysis_->GetPointsToSet(instruction); - return points_to_set.ForEachElementWithStatus( + const InstructionValueSet& value_set_tree = + dataflow_analysis_->GetInstructionValueSet(instruction); + return value_set_tree.ForEachElementWithStatus( [&shape_layout, this, priority, user]( const ShapeIndex& index, - const PointsToSet::BufferList& buffers) -> absl::Status { + const HloValueSet& value_set) -> absl::Status { const auto& subshape = ShapeUtil::GetSubshape(shape_layout.shape(), index); if (ShapeUtil::IsLeafIndex(shape_layout.shape(), index) && subshape.has_layout()) { - for (const LogicalBuffer* buffer : buffers) { + for (const HloValue* buffer : value_set.values()) { if (buffer->shape().IsArray() && (buffer->instruction()->opcode() != HloOpcode::kReduce || !buffer->instruction()->shape().IsTuple())) { @@ -2493,16 +2507,14 @@ absl::Status LayoutAssignment::PropagateOperandConstraintToResultForCustomCall( continue; } ShapeIndex shape_index = output_operand_pair.first; - if (!points_to_analysis_->InstructionDefinesBufferAtIndex(user, - shape_index)) { + if (!dataflow_analysis_->ValueIsDefinedAt(user, shape_index)) { return absl::OkStatus(); } - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); return SetBufferLayout( - operand_constraint.shape_layout().layout(), *buffer, + operand_constraint.shape_layout().layout(), buffer, /*mandatory=*/OperandLayoutAlwaysPropagateForward(user), /*dfs=*/InstructionShouldPropagateDepthFirst(*user), operand_constraint.priority()); @@ -2552,11 +2564,10 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( user->operand_count() == 1 ? ShapeIndex() : ShapeIndex({operand_constraint.operand_no()}); - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); ABSL_RETURN_IF_ERROR(SetBufferLayout(operand_constraint.shape_layout().layout(), - *buffer, + buffer, /*mandatory=*/true, /*dfs=*/true)); } @@ -2631,22 +2642,20 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( operand->shape().dimensions().size()) { return absl::OkStatus(); } - if (!points_to_analysis_->InstructionDefinesBufferAtIndex( - user, shape_index)) { + if (!dataflow_analysis_->ValueIsDefinedAt(user, shape_index)) { return absl::OkStatus(); } // TODO(b/67641796): Are there cases except fusion that use this code // path? - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); // If we already have a constraint for the buffer it was assigned but // hasn't propagated yet. This can happen with diamond-shaped graphs // where one path is first evaluated in depth-first order (we're here) // and the other path is propagated later. We don't set the layout // here as it will always be overwritten later. ABSL_RETURN_IF_ERROR(SetBufferLayout( - operand_constraint.shape_layout().layout(), *buffer, + operand_constraint.shape_layout().layout(), buffer, /*mandatory=*/true, /*dfs=*/true, operand_constraint.priority())); return absl::OkStatus(); })); @@ -2662,19 +2671,17 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( if (subshape.dimensions().size() <= 1) { return absl::OkStatus(); } - if (!points_to_analysis_->InstructionDefinesBufferAtIndex( - user, shape_index)) { + if (!dataflow_analysis_->ValueIsDefinedAt(user, shape_index)) { return absl::OkStatus(); } - ABSL_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - points_to_analysis_->GetBufferDefinedAt(user, shape_index)); + const HloValue& buffer = + dataflow_analysis_->GetValueDefinedAt(user, shape_index); std::unique_ptr layout = ChooseOutputLayoutFromOperandLayout( operand_constraint.shape_layout().layout(), user, operand_constraint.operand_no()); if (layout != nullptr) { ABSL_RETURN_IF_ERROR(SetBufferLayout( - *layout, *buffer, + *layout, buffer, /*mandatory=*/OperandLayoutAlwaysPropagateForward(user), /*dfs=*/InstructionShouldPropagateDepthFirst(*user), operand_constraint.priority())); @@ -2687,7 +2694,7 @@ absl::Status LayoutAssignment::PropagateOperandConstraint( absl::Status LayoutAssignment::PropagateBufferConstraintToOperands( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) { - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); const HloInstruction* instruction = buffer.instruction(); if (IsAtMostRank1(instruction->shape())) { return absl::OkStatus(); @@ -2768,7 +2775,7 @@ absl::Status LayoutAssignment::PropagateBufferConstraint( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) { // Only propagate array layouts. - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); if (!buffer.IsArray()) { return absl::OkStatus(); } @@ -2782,13 +2789,12 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToUses( LayoutConstraints* constraints) { VLOG(5) << "PropagateBufferConstraintToUses: " << buffer_constraint.ToString(); - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); TF_RET_CHECK(buffer.IsArray()); // Propagate the layout to all array uses of the logical buffer. This skips // uses of the buffer where the buffer is the element of a tuple. - for (const auto& user_operand_no : - GetArrayUsesOfBuffer(points_to_analysis_->GetBufferAliases(buffer))) { + for (const auto& user_operand_no : GetArrayUsesOfBuffer(buffer)) { const HloInstruction* user = user_operand_no.first; int64_t operand_no = user_operand_no.second; // Only add an operand constraint if the user does not forward the buffer @@ -2825,11 +2831,12 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToUses( ShapeIndex used_index = buffer.index(); used_index.push_front(index); - ABSL_ASSIGN_OR_RETURN(auto buffer, points_to_analysis_->GetBufferDefinedAt( - inputs, used_index)); - - ABSL_RETURN_IF_ERROR(SetBufferLayout(buffer_constraint.layout(), *buffer, - /*mandatory=*/false)); + for (const HloValue* input_buffer : + dataflow_analysis_->GetValueSet(inputs, used_index).values()) { + ABSL_RETURN_IF_ERROR(SetBufferLayout(buffer_constraint.layout(), + *input_buffer, + /*mandatory=*/false)); + } } } @@ -2861,30 +2868,30 @@ absl::Status LayoutAssignment::PropagateResultConstraint( absl::StatusOr LayoutAssignment::InferArrayLayout( const HloInstruction* instruction, const ShapeIndex& index) { const auto& source_buffers = - points_to_analysis_->GetPointsToSet(instruction).element(index); + dataflow_analysis_->GetValueSet(instruction, index).values(); TF_RET_CHECK(!source_buffers.empty()); - // Verify the layout is the same for every LogicalBuffer which this location + // Verify the layout is the same for every HloValue which this location // ('instruction' and 'index') points to. const Layout* first_buffer_layout = nullptr; - for (const LogicalBuffer* source_buffer : source_buffers) { + for (const HloValue* source_buffer : source_buffers) { VLOG(5) << "Logical buffer: " << source_buffer->ToString() << "\n"; auto* source_buffer_constraint = GetBufferLayoutConstraint(*source_buffer); if (source_buffer_constraint == nullptr) { // This should not happen because we've assigned layouts to all // instructions preceding this one. - return Internal("LogicalBuffer %s does not have a layout", + return Internal("HloValue %s does not have a layout", source_buffer->ToString()); } if (first_buffer_layout == nullptr) { first_buffer_layout = &source_buffer_constraint->layout(); } else if (!Layout::Equal().MinorToMajorOnly()( - source_buffer->shape().layout(), *first_buffer_layout)) { - // The points-to set is ambiguous for this index and the different source + source_buffer_constraint->layout(), *first_buffer_layout)) { + // The value set is ambiguous for this index and the different source // buffers have different layouts. This case is possible in valid XLA // computations because we do not propagate BufferLayoutConstraints to all - // LogicalBuffers which may alias the constrained LogicalBuffer at some + // HloValues which may alias the constrained HloValue at some // point in the computation. return FailedPrecondition( "Array at index {%s} in instruction %s aliases buffers %s " @@ -2967,23 +2974,26 @@ absl::Status LayoutAssignment::AssignLayouts(LayoutConstraints& constraints) { // by the respective BufferLayoutConstraints. Any array shapes in the output // of the instruction which are not defined by the instruction (eg, array // elements in a Tuple instruction) will be assigned below via inference. - for (const LogicalBuffer* buffer : - points_to_analysis_->GetBuffersDefinedByInstruction(instruction)) { - if (!buffer->shape().IsArray()) { - continue; - } - TF_RET_CHECK(buffer->instruction() == instruction); - auto* buffer_layout_constraint = GetBufferLayoutConstraint(*buffer); - TF_RET_CHECK(buffer_layout_constraint != nullptr); - if (instruction->opcode() == HloOpcode::kConstant) { - // For constants, we also need to change the layout of the internal - // literal. - instruction->RelayoutConstant(buffer_layout_constraint->layout(), - buffer->index()); - } else { - Shape* buffer_subshape = ShapeUtil::GetMutableSubshape( - instruction->mutable_shape(), buffer->index()); - *buffer_subshape->mutable_layout() = buffer_layout_constraint->layout(); + for (const auto& [index, value_set] : + dataflow_analysis_->GetInstructionValueSet(instruction)) { + for (const HloValue* buffer : value_set.values()) { + if (!buffer->shape().IsArray() || + buffer->instruction() != instruction) { + continue; + } + auto* buffer_layout_constraint = GetBufferLayoutConstraint(*buffer); + TF_RET_CHECK(buffer_layout_constraint != nullptr); + if (instruction->opcode() == HloOpcode::kConstant) { + // For constants, we also need to change the layout of the internal + // literal. + instruction->RelayoutConstant(buffer_layout_constraint->layout(), + buffer->index()); + } else { + Shape* buffer_subshape = ShapeUtil::GetMutableSubshape( + instruction->mutable_shape(), buffer->index()); + *buffer_subshape->mutable_layout() = + buffer_layout_constraint->layout(); + } } } @@ -3261,10 +3271,10 @@ absl::Status LayoutAssignment::ClearComputationLayouts( void LayoutAssignment::InitUnconstrainedBuffers(HloComputation* computation) { for (HloInstruction* inst : computation->instructions()) { - points_to_analysis_->GetPointsToSet(inst).ForEachElement( - [&](const ShapeIndex&, const PointsToSet::BufferList& buffers) { - for (const LogicalBuffer* buffer : buffers) { - // The points to analysis is computed per module, restrict + dataflow_analysis_->GetInstructionValueSet(inst).ForEachElement( + [&](const ShapeIndex&, const HloValueSet& value_set) { + for (const HloValue* buffer : value_set.values()) { + // The dataflow analysis is computed per module, restrict // constraints to array buffers in this computation. if (buffer->IsArray() && buffer->instruction()->parent() == computation) { @@ -3311,13 +3321,9 @@ absl::Status LayoutAssignment::AddCustomCallConstraints( } void LayoutAssignment::RecordUnconstrainedLayoutInstructions() { - for (LogicalBuffer::Id buffer_id : unconstrained_buffer_ids_) { - VLOG(5) - << "unconstrained instruction:" - << points_to_analysis_->GetBuffer(buffer_id).instruction()->ToString() - << "\n"; + for (HloValue::Id buffer_id : unconstrained_buffer_ids_) { unconstrained_layout_instructions_.insert( - points_to_analysis_->GetBuffer(buffer_id).instruction()); + dataflow_analysis_->GetValue(buffer_id).instruction()); } } @@ -3328,9 +3334,9 @@ absl::Status LayoutAssignment::AssignLayoutsToUnconstrainedBuffers( // Arbitrarily pick the first unconstrained buffer and assign a layout. // Constants use their literal's layout; other buffers get a default layout. - // unconstrained_buffer_ids_ has a stable sort based on LogicalBuffer::Id. - const LogicalBuffer& buffer = - points_to_analysis_->GetBuffer(*unconstrained_buffer_ids_.begin()); + // unconstrained_buffer_ids_ has a stable sort based on HloValue::Id. + const HloValue& buffer = + dataflow_analysis_->GetValue(*unconstrained_buffer_ids_.begin()); const HloInstruction* instruction = buffer.instruction(); Layout new_layout = instruction->opcode() == HloOpcode::kConstant @@ -3598,8 +3604,14 @@ absl::Status LayoutAssignment::VerifyEntryComputationLayout( absl::StatusOr> LayoutAssignment::SetupPropagation( HloModule* module, const absl::flat_hash_set& execution_threads) { - ABSL_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); - SetPointsToAnalysis(std::move(points_to_analysis)); + ABSL_ASSIGN_OR_RETURN(auto dataflow_analysis, + HloDataflowAnalysis::Run( + *module, /*ssa_form=*/false, + /*bitcast_defines_value=*/false, execution_threads, + /*propagate_through_calls=*/false, + /*precompute_uses=*/std::nullopt, + /*propagate_through_control_flow=*/false)); + SetDataflowAnalysis(std::move(dataflow_analysis)); auto computations_to_work = module->MakeNonfusionComputations(execution_threads); // If the reverse_computation_order_ flag is set, reverse the ordering of @@ -3962,8 +3974,8 @@ absl::Status LayoutAssignment::ClearPreviousPassSideEffects( unconstrained_layout_instructions_.clear(); unconstrained_buffer_ids_.clear(); buffer_constraints_.Clear(); - // buffer_sets_cache_ stays: it depends only on points_to_analysis_, which - // SetPointsToAnalysis owns. + // buffer_sets_cache_ stays: it depends only on dataflow_analysis_, which + // SetDataflowAnalysis owns. return absl::OkStatus(); } absl::Status LayoutAssignment::AddCopyForOperand(HloInstruction* instruction, diff --git a/third_party/xla/xla/service/layout_assignment.h b/third_party/xla/xla/service/layout_assignment.h index abcfe3eb5a23a6..934b35780a8590 100644 --- a/third_party/xla/xla/service/layout_assignment.h +++ b/third_party/xla/xla/service/layout_assignment.h @@ -33,7 +33,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" -#include "xla/hlo/analysis/tuple_points_to_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" @@ -43,7 +43,7 @@ limitations under the License. #include "xla/layout_util.h" #include "xla/map_util.h" #include "xla/service/computation_layout.h" -#include "xla/service/logical_buffer.h" +#include "xla/service/hlo_value.h" #include "xla/shape.h" #include "xla/shape_layout.h" #include "xla/shape_util.h" @@ -137,12 +137,12 @@ std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint); // or induced this buffer layout constraint. class BufferLayoutConstraint : public LayoutConstraint { public: - // Constructs a BufferLayoutConstraint for the given LogicalBuffer. - BufferLayoutConstraint(const Layout& layout, const LogicalBuffer& buffer, + // Constructs a BufferLayoutConstraint for the given HloValue. + BufferLayoutConstraint(const Layout& layout, const HloValue& buffer, bool mandatory, bool dfs, int64_t priority); - // Returns the constrained LogicalBuffer. - const LogicalBuffer& buffer() const { return *buffer_; } + // Returns the constrained HloValue. + const HloValue& buffer() const { return *buffer_; } // Returns the target Layout for the buffer. const Layout& layout() const { return layout_[0]; } @@ -159,8 +159,8 @@ class BufferLayoutConstraint : public LayoutConstraint { // The constrained layout(s) for the buffer. absl::InlinedVector layout_; - // The LogicalBuffer being constrained. - const LogicalBuffer* buffer_; + // The HloValue being constrained. + const HloValue* buffer_; // The consumer instruction that induced this constraint, if any. const HloInstruction* from_user_ = nullptr; @@ -400,8 +400,8 @@ class LayoutAssignment : public HloModulePass { ChannelLayoutConstraints* channel_constraints = nullptr, bool reverse_computation_order = false); ~LayoutAssignment() override {} - const TuplePointsToAnalysis& points_to_analysis() const { - return *points_to_analysis_; + const HloDataflowAnalysis& dataflow_analysis() const { + return *dataflow_analysis_; } absl::string_view name() const override { return "layout-assignment"; } @@ -574,14 +574,12 @@ class LayoutAssignment : public HloModulePass { // Add a constraint on the layout of a LogicalBuffer, the layout of the // operand of the instruction, or the layout of the result of the computation, // respectively. - absl::Status SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer, + absl::Status SetBufferLayout(const Layout& layout, const HloValue& buffer, bool mandatory = true, bool dfs = true) { return SetBufferLayout(layout, buffer, mandatory, dfs, current_priority_); } - absl::Status SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer, bool mandatory, - bool dfs, int64_t priority, + absl::Status SetBufferLayout(const Layout& layout, const HloValue& buffer, + bool mandatory, bool dfs, int64_t priority, const HloInstruction* from_user = nullptr); absl::Status SetOperandLayout(const Shape& shape_with_layout, const HloInstruction* instruction, @@ -647,7 +645,7 @@ class LayoutAssignment : public HloModulePass { const ComputationLayoutConstraint& layout_constraint, LayoutConstraints* constraints); - virtual Layout GetUnconstrainedLayout(const LogicalBuffer& buffer) { + virtual Layout GetUnconstrainedLayout(const HloValue& buffer) { Layout layout = LayoutUtil::GetDefaultLayoutForShape(buffer.shape()); if (buffer.shape().has_layout()) { layout.set_memory_space(buffer.shape().layout().memory_space()); @@ -662,13 +660,14 @@ class LayoutAssignment : public HloModulePass { absl::Status PropagateUnconstraintedBuffers(LayoutConstraints* constraints); const BufferLayoutConstraint* GetBufferLayoutConstraint( - const LogicalBuffer& buffer) const; + const HloValue& buffer) const; absl::StatusOr GetInstructionBufferLayoutConstraint(const HloInstruction* instruction) const; + using BufferSet = absl::flat_hash_set; // Find a bufferset in the bufferset cache. This is useful since we can // currently create the flattened buffer set for the same instruction many // times, which is often slow. - PointsToSet::BufferSet* GetBufferSet(const HloInstruction* instruction) const; + BufferSet* GetBufferSet(const HloInstruction* instruction) const; // Similar to above, but returns true only if all buffers associated with that // operand are forwarded. bool AllOperandBuffersForwarded(const HloInstruction* instruction, @@ -871,8 +870,8 @@ class LayoutAssignment : public HloModulePass { // Constrains layouts for custom calls that have specific layout requirements. absl::Status AddCustomCallConstraints(LayoutConstraints* constraints); - // Initializes unconstrained_buffer_ids_ with all array-shaped logical buffers - // in the given computation. + // Initializes unconstrained_buffer_ids_ with all array HloValues in the + // given computation. void InitUnconstrainedBuffers(HloComputation* computation); // Records instructions that lack layout constraints before applying default @@ -1023,12 +1022,12 @@ class LayoutAssignment : public HloModulePass { // Adds constraints related to host Send/Recv instructions. absl::Status BuildHostChannelConstraints(HloComputation* computation); - // Replaces points_to_analysis_ and drops the buffer sets memoized from the + // Replaces dataflow_analysis_ and drops the buffer sets memoized from the // previous analysis. - void SetPointsToAnalysis(std::unique_ptr analysis); + void SetDataflowAnalysis(std::unique_ptr analysis); - // Module points to analysis that can be updated for cloned computations. - std::unique_ptr points_to_analysis_; + // Module dataflow analysis that can be updated for cloned computations. + std::unique_ptr dataflow_analysis_; // The set of HLO instructions which lacked any layout constraint, thus // receiving propagated default layouts. @@ -1072,25 +1071,24 @@ class LayoutAssignment : public HloModulePass { ChannelLayoutConstraints host_channel_constraints_; // Array-shaped buffers which have not yet been constrained, in id order. - absl::btree_set unconstrained_buffer_ids_; + absl::btree_set unconstrained_buffer_ids_; - // Buffer sets of points_to_analysis_ memoized by GetBufferSet. The entries + // Buffer sets of dataflow_analysis_ memoized by GetBufferSet. The entries // point at buffers owned by the analysis, so the cache is only valid for the - // analysis it was filled from: SetPointsToAnalysis, the only place the + // analysis it was filled from: SetDataflowAnalysis, the only place the // analysis changes, clears it, and GetBufferSet checks that pairing. - mutable absl::flat_hash_map> + mutable absl::flat_hash_map> buffer_sets_cache_; - const TuplePointsToAnalysis* buffer_sets_cache_analysis_ = nullptr; + const HloDataflowAnalysis* buffer_sets_cache_analysis_ = nullptr; - // Buffer layout constraints stored densely by LogicalBuffer::Id, which the - // points-to analysis assigns sequentially. Clear() invalidates every entry + // Buffer layout constraints stored densely by HloValue::Id, which the + // dataflow analysis assigns sequentially. Clear() invalidates every entry // in O(1) by bumping a generation, so no table is destroyed and rebuilt per // propagation round. class BufferConstraintTable { public: // The live constraint of `id`, or nullptr. - BufferLayoutConstraint* Find(LogicalBuffer::Id id) const { + BufferLayoutConstraint* Find(HloValue::Id id) const { const size_t slot = static_cast(id); return slot < constraints_.size() && generations_[slot] == generation_ ? constraints_[slot].get() @@ -1099,8 +1097,7 @@ class LayoutAssignment : public HloModulePass { // Makes `constraint` the live constraint of `id`. BufferLayoutConstraint* Insert( - LogicalBuffer::Id id, - std::unique_ptr constraint) { + HloValue::Id id, std::unique_ptr constraint) { const size_t slot = static_cast(id); if (slot >= constraints_.size()) { constraints_.resize(slot + 1); diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc index b7268e3234e460..c49e7d4bf1fe2d 100644 --- a/third_party/xla/xla/service/layout_assignment_test.cc +++ b/third_party/xla/xla/service/layout_assignment_test.cc @@ -43,8 +43,8 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/computation_layout.h" +#include "xla/service/hlo_value.h" #include "xla/service/hlo_verifier.h" -#include "xla/service/logical_buffer.h" #include "xla/service/pattern_matcher.h" #include "xla/shape.h" #include "xla/shape_layout.h" @@ -525,7 +525,7 @@ class OperandsMustBeTheSameLayoutAssignment : public LayoutAssignment { absl::Status PropagateBufferConstraint( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) override { - const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloValue& buffer = buffer_constraint.buffer(); const HloInstruction* instruction = buffer.instruction(); // Force the operands' layout to the output layout. diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index 2d863d1a66fc41..ba5e63312be9f5 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -1093,7 +1093,7 @@ cc_library( name = "command_buffer", srcs = ["command_buffer.cc"], hdrs = ["command_buffer.h"], - visibility = [":internal"], + visibility = internal_visibility([":internal"]), deps = [ ":bit_pattern", ":device_address", diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc index 39f1ff7fa0afe3..c8935d0a0308bb 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc +++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.cc @@ -2966,10 +2966,11 @@ absl::Status MIOpenSupport::GetConvolveRunners( } std::vector profile_results; - if (!GetMIOpenConvolveAlgorithms( + if (!GetMIOpenConvolveAlgorithmsInternal( kind, input_type, output_type, stream, input_descriptor, input_data, filter_descriptor, filter_data, output_descriptor, output_data, - convolution_descriptor, scratch_allocator, &profile_results)) + convolution_descriptor, scratch_allocator, &profile_results, + use_fallback)) return absl::InternalError("GetMIOpenConvolveAlgorithms failure"); for (const auto& profile_result : profile_results) { @@ -3092,13 +3093,33 @@ bool MIOpenSupport::GetMIOpenConvolveAlgorithms( const dnn::ConvolutionDescriptor& convolution_descriptor, ScratchAllocator* scratch_allocator, std::vector* out_algorithms) { + return GetMIOpenConvolveAlgorithmsInternal( + kind, input_type, output_type, stream, input_descriptor, input_data, + filter_descriptor, filter_data, output_descriptor, output_data, + convolution_descriptor, scratch_allocator, out_algorithms, + /*use_fallback=*/false); +} + +bool MIOpenSupport::GetMIOpenConvolveAlgorithmsInternal( + dnn::ConvolutionKind kind, dnn::DataType input_type, + dnn::DataType output_type, Stream* stream, + const dnn::BatchDescriptor& input_descriptor, DeviceAddressBase input_data, + const dnn::FilterDescriptor& filter_descriptor, + DeviceAddressBase filter_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceAddressBase output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + ScratchAllocator* scratch_allocator, + std::vector* out_algorithms, bool use_fallback) { // TODO(rocm): Create handles only once and reuse them between the methods - if (!PopulateMIOpenFindDb(kind, input_type, output_type, stream, - input_descriptor, input_data, filter_descriptor, - filter_data, output_descriptor, output_data, - convolution_descriptor, scratch_allocator) - .ok()) { - return false; + if (!use_fallback) { + if (!PopulateMIOpenFindDb(kind, input_type, output_type, stream, + input_descriptor, input_data, filter_descriptor, + filter_data, output_descriptor, output_data, + convolution_descriptor, scratch_allocator) + .ok()) { + return false; + } } return GetMIOpenConvolveAlgorithmsImmediateMode( kind, input_type, output_type, stream, input_descriptor, diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h index c293839253e1f9..e2f3044166096f 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_dnn.h @@ -256,6 +256,19 @@ class MIOpenSupport : public dnn::DnnSupport { ScratchAllocator* scratch_allocator, std::vector* out_algorithms) override; + bool GetMIOpenConvolveAlgorithmsInternal( + dnn::ConvolutionKind kind, dnn::DataType input_type, + dnn::DataType output_type, Stream* stream, + const dnn::BatchDescriptor& input_descriptor, + DeviceAddressBase input_data, + const dnn::FilterDescriptor& filter_descriptor, + DeviceAddressBase filter_data, + const dnn::BatchDescriptor& output_descriptor, + DeviceAddressBase output_data, + const dnn::ConvolutionDescriptor& convolution_descriptor, + ScratchAllocator* scratch_allocator, + std::vector* out_algorithms, bool use_fallback); + absl::Status GetMIOpenConvolveAlgorithmsImmediateMode( dnn::ConvolutionKind kind, dnn::DataType input_type, dnn::DataType output_type, Stream* stream, diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc index 89f7202922dd19..a98bd795160457 100644 --- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc +++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc @@ -373,9 +373,6 @@ absl::StatusOr FetchAndLogOutput( const std::vector>>& output_buffers, ModuleOutputMode module_output_mode, bool log_output) { CHECK(!output_buffers.empty()); - absl::Mutex mu; - absl::Status status; - size_t num_pending_transfers = 0; bool device_0_is_local = false; for (PjRtDevice* device : GetLocalDevices(client)) { if (device->id() == 0) { @@ -383,67 +380,81 @@ absl::StatusOr FetchAndLogOutput( } } - if (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_0_is_local) { - num_pending_transfers = output_buffers[0].size(); - } else if (module_output_mode == ModuleOutputMode::kReturnOutputs) { - for (const auto& bs : output_buffers) { - num_pending_transfers += bs.size(); - } - } - PerDeviceLiteralVecType outputs; - for (int i = 0; i < output_buffers.size(); ++i) { - if (output_buffers[i].empty()) { - continue; - } - const int device_id = output_buffers[i][0]->device()->id(); - std::vector& output_slice = outputs[device_id]; - if (module_output_mode == ModuleOutputMode::kReturnOutputs || - (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_id == 0)) { - output_slice.reserve(output_buffers[i].size()); - for (const auto& buffer : output_buffers[i]) { - TF_RET_CHECK(buffer->device() == output_buffers[i][0]->device()) - << "All outputs from a given vector of outputs should be for the " - "same device"; - ABSL_ASSIGN_OR_RETURN(auto logical_shape, buffer->logical_on_device_shape()); - output_slice.emplace_back( - ShapeUtil::DeviceShapeToHostShape(logical_shape)); - buffer->ToLiteral(&output_slice.back()).OnReady([&](absl::Status s) { - absl::MutexLock lock(mu); - --num_pending_transfers; - status.Update(s); - }); + absl::Mutex mu; + absl::Status status; + size_t num_pending_transfers = 0; + + absl::Status issue_status = [&]() -> absl::Status { + for (int i = 0; i < output_buffers.size(); ++i) { + if (output_buffers[i].empty()) { + continue; } - } else { - for (const auto& buffer : output_buffers[i]) { - TF_RET_CHECK(buffer->device() == output_buffers[i][0]->device()) - << "All outputs from a given vector of outputs should be for the " - "same device"; - ABSL_RETURN_IF_ERROR(buffer->GetReadyFuture().Await()); + const int device_id = output_buffers[i][0]->device()->id(); + std::vector& output_slice = outputs[device_id]; + if (module_output_mode == ModuleOutputMode::kReturnOutputs || + (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && + device_id == 0)) { + output_slice.reserve(output_buffers[i].size()); + for (const auto& buffer : output_buffers[i]) { + if (buffer->device() != output_buffers[i][0]->device()) { + return absl::InternalError( + "All outputs from a given vector of outputs should be for the " + "same device"); + } + ABSL_ASSIGN_OR_RETURN(auto logical_shape, + buffer->logical_on_device_shape()); + output_slice.emplace_back( + ShapeUtil::DeviceShapeToHostShape(logical_shape)); + { + absl::MutexLock lock(mu); + ++num_pending_transfers; + } + buffer->ToLiteral(&output_slice.back()).OnReady([&](absl::Status s) { + absl::MutexLock lock(mu); + --num_pending_transfers; + status.Update(s); + }); + } + } else { + for (const auto& buffer : output_buffers[i]) { + if (buffer->device() != output_buffers[i][0]->device()) { + return absl::InternalError( + "All outputs from a given vector of outputs should be for the " + "same device"); + } + ABSL_RETURN_IF_ERROR(buffer->GetReadyFuture().Await()); + } } } - } - if (module_output_mode == ModuleOutputMode::kReturnOutputs || - (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_0_is_local)) { - auto cond = [&]() { return !status.ok() || num_pending_transfers == 0; }; + return absl::OkStatus(); + }(); + + // The ToLiteral callbacks above reference `outputs`, `mu`, `status` and + // `num_pending_transfers`, so every issued transfer must complete before this + // function returns, including when a transfer fails or issuing one fails. + { + auto all_transfers_done = [&]() { return num_pending_transfers == 0; }; absl::MutexLock lock(mu); - mu.Await(absl::Condition(&cond)); + mu.Await(absl::Condition(&all_transfers_done)); + ABSL_RETURN_IF_ERROR(issue_status); ABSL_RETURN_IF_ERROR(status); - if (log_output) { - for (const PjRtDevice* device : GetLocalDevices(client)) { - int device_id = device->id(); - if (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && - device_id != 0) { - continue; - } - LOG(INFO) << "Outputs for device_id: " << device_id; - const std::vector& output_slice = outputs[device_id]; - for (int i = 0; i < output_slice.size(); ++i) { - LOG(INFO) << "output[" << i << "]: " << output_slice[i].ToString(); - } + } + + if (log_output && + (module_output_mode == ModuleOutputMode::kReturnOutputs || + (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && + device_0_is_local))) { + for (const PjRtDevice* device : GetLocalDevices(client)) { + int device_id = device->id(); + if (module_output_mode == ModuleOutputMode::kReturnDevice0Outputs && + device_id != 0) { + continue; + } + LOG(INFO) << "Outputs for device_id: " << device_id; + const std::vector& output_slice = outputs[device_id]; + for (int i = 0; i < output_slice.size(); ++i) { + LOG(INFO) << "output[" << i << "]: " << output_slice[i].ToString(); } } } diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index a23060f7f1ee10..4a56ebe77e5e97 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -94,6 +94,7 @@ message DebugOptions { DYNAMIC_SLICE_FUSION = 9; CONVOLUTION = 11; HOST_EXECUTE = 12; + COLLECTIVES_KERNEL = 13; reserved "DYNAMIC_SLICE_COPY_FUSION"; reserved 10; From f235c9edfbeafa95b663d0238bfaac4c683003d7 Mon Sep 17 00:00:00 2001 From: Evan Brown Date: Tue, 29 Sep 2026 04:45:15 -0700 Subject: [PATCH 08/25] Automated Code Change PiperOrigin-RevId: 990223584 --- tensorflow/core/common_runtime/rendezvous_util.cc | 4 ++-- tensorflow/core/distributed_runtime/graph_mgr.cc | 2 -- .../xla/xla/tsl/platform/cloud/gcs_file_system_test.cc | 9 ++++++--- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/tensorflow/core/common_runtime/rendezvous_util.cc b/tensorflow/core/common_runtime/rendezvous_util.cc index 5e28d6ac5d70bc..dbfabfb9c1dcf1 100644 --- a/tensorflow/core/common_runtime/rendezvous_util.cc +++ b/tensorflow/core/common_runtime/rendezvous_util.cc @@ -75,14 +75,14 @@ void RecvOutputsFromRendezvousAsync( keys.size(), "; alloc_attrs.size() = ", alloc_attrs.size()))); } - received_tensors->reserve(keys.size()); + received_tensors->resize(keys.size()); std::vector> arguments; + arguments.reserve(keys.size()); for (int i = 0; i < keys.size(); ++i) { Rendezvous::ParsedKey parsed; absl::Status s = Rendezvous::ParseKey(keys[i], &parsed); - received_tensors->push_back(Tensor()); if (!s.ok()) { done(s); return; diff --git a/tensorflow/core/distributed_runtime/graph_mgr.cc b/tensorflow/core/distributed_runtime/graph_mgr.cc index d50a0f1d67f6b6..75375e0f5f47b3 100644 --- a/tensorflow/core/distributed_runtime/graph_mgr.cc +++ b/tensorflow/core/distributed_runtime/graph_mgr.cc @@ -427,10 +427,8 @@ void GraphMgr::RecvOutputsAsync(const int64_t step_id, NamedTensors* out, std::vector keys; std::vector* received_keys = new std::vector; keys.reserve(out->size()); - received_keys->reserve(out->size()); for (const auto& p : *out) { keys.push_back(p.first); - received_keys->push_back(p.second); } RecvOutputsFromRendezvousAsync( rendezvous, nullptr, {}, keys, received_keys, diff --git a/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc b/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc index 76f2d24379504f..4e0fc2e020e358 100644 --- a/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc +++ b/third_party/xla/xla/tsl/platform/cloud/gcs_file_system_test.cc @@ -634,7 +634,8 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache) { std::unique_ptr( new FakeHttpRequestFactory(&requests)), std::unique_ptr(new FakeZoneProvider), 9 /* block size */, - 18 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */, + 1 << 20 /* max bytes: large enough that nothing is evicted */, + 0 /* max staleness */, 3600 /* stat cache max age */, 0 /* stat cache max entries */, 0 /* matching paths cache max age */, 0 /* matching paths cache max entries */, kTestRetryConfig, kTestTimeoutConfig, *kAllowedLocationsDefault, @@ -723,7 +724,8 @@ TEST(GcsFileSystemTest, NewRandomAccessFile_WithBlockCache_Flush) { std::unique_ptr( new FakeHttpRequestFactory(&requests)), std::unique_ptr(new FakeZoneProvider), 9 /* block size */, - 18 /* max bytes */, 0 /* max staleness */, 3600 /* stat cache max age */, + 1 << 20 /* max bytes: large enough that nothing is evicted */, + 0 /* max staleness */, 3600 /* stat cache max age */, 0 /* stat cache max entries */, 0 /* matching paths cache max age */, 0 /* matching paths cache max entries */, kTestRetryConfig, kTestTimeoutConfig, *kAllowedLocationsDefault, @@ -843,7 +845,8 @@ TEST(GcsFileSystemTest, std::unique_ptr( new FakeHttpRequestFactory(&requests)), std::unique_ptr(new FakeZoneProvider), 9 /* block size */, - 18 /* max bytes */, 0 /* max staleness */, 0 /* stat cache max age */, + 1 << 20 /* max bytes: large enough that nothing is evicted */, + 0 /* max staleness */, 0 /* stat cache max age */, 0 /* stat cache max entries */, 0 /* matching paths cache max age */, 0 /* matching paths cache max entries */, kTestRetryConfig, kTestTimeoutConfig, *kAllowedLocationsDefault, From 19f27949612f39fe0a4206993761b381d5b773b4 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Tue, 29 Sep 2026 05:13:46 -0700 Subject: [PATCH 09/25] Enable libnvjitlink unconditionally in XLA. Previously, libnvjitlink support was gated behind the `enable_libnvjitlink_support` build flag (which defaulted to `false` in OSS) and split across `nvjitlink_impl` and `nvjitlink_stub` targets. Because all supported CUDA versions now include libnvjitlink, the conditional build graph is no longer needed. Compile xla/stream_executor/cuda/nvjitlink.cc unconditionally in //xla/stream_executor/cuda:nvjitlink and remove the `enable_libnvjitlink_support` build flag, `libnvjitlink_support_enabled` config setting, and xla/stream_executor/cuda/nvjitlink_stub.cc. Keep //xla/stream_executor/cuda:nvjitlink_support only to disable libnvjitlink under MemorySanitizer. To support OSS wheel builds that dynamically load CUDA libraries via stubs, add //xla/tsl/cuda:nvjitlink (backed by xla/tsl/cuda/nvjitlink_stub.cc and `GetNvJitLinkDsoHandle()` in xla/tsl/platform/default/dso_loader.cc). In `kAuto` compilation mode, check `GetNvJitLinkVersion()` before `LoadedNvJitLinkHasKnownIssues()` so stub builds without `libnvJitLink.so` at runtime fall back cleanly to `ptxas` and `nvlink`. PiperOrigin-RevId: 990235320 --- third_party/xla/xla/BUILD | 1 - third_party/xla/xla/debug_options_flags.cc | 3 +- .../xla/xla/stream_executor/cuda/BUILD | 109 ++++----------- .../cuda/assemble_compilation_provider.cc | 12 ++ .../cuda/{nvjitlink_impl.cc => nvjitlink.cc} | 25 +++- .../cuda/nvjitlink_known_issues_test.cc | 19 ++- .../stream_executor/cuda/nvjitlink_stub.cc | 40 ------ .../stream_executor/cuda/nvjitlink_support.cc | 2 +- .../stream_executor/cuda/nvjitlink_test.cc | 130 +++++++++++++++--- third_party/xla/xla/tsl/cuda/BUILD.bazel | 49 +++++++ .../xla/xla/tsl/cuda/nvjitlink.symbols | 16 +++ .../xla/xla/tsl/cuda/nvjitlink_stub.cc | 73 ++++++++++ .../xla/tsl/platform/default/dso_loader.cc | 12 +- .../xla/xla/tsl/platform/default/dso_loader.h | 6 + 14 files changed, 343 insertions(+), 154 deletions(-) rename third_party/xla/xla/stream_executor/cuda/{nvjitlink_impl.cc => nvjitlink.cc} (91%) delete mode 100644 third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc create mode 100644 third_party/xla/xla/tsl/cuda/nvjitlink.symbols create mode 100644 third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index b720b59902e79c..406a441dc69dcc 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1353,7 +1353,6 @@ cc_library( "//xla/backends/autotuner:backends_proto_cc", "//xla/hlo/pass:hlo_pass_filter", "//xla/service:collective_utils", - "//xla/stream_executor/cuda:nvjitlink_support", "//xla/stream_executor/cuda:ptx_compiler_support", "//xla/tsl/platform:logging", "//xla/tsl/util:command_line_flags", diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index a86920cee08c4d..4ff61a22dd4f24 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -54,7 +54,6 @@ limitations under the License. #include "xla/hlo/pass/hlo_pass_filter.h" #include "xla/parse_flags_from_env.h" #include "xla/service/collective_utils.h" -#include "xla/stream_executor/cuda/nvjitlink_support.h" #include "xla/stream_executor/cuda/ptx_compiler_support.h" #include "xla/tsl/platform/logging.h" // IWYU pragma: keep #include "xla/tsl/util/command_line_flags.h" @@ -2959,7 +2958,7 @@ void MakeDebugOptionsFlags(std::vector* flag_list, : DebugOptions::LIB_NV_JIT_LINK_MODE_DISABLED); return true; }, - stream_executor::IsLibNvJitLinkSupported(), + /*default_value_for_display=*/true, "Use libnvjitlink for PTX-to-GPU-assembly compilation instead of " "calling ptxas.")); flag_list->push_back(tsl::Flag( diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD index d23960e276d69b..5d69ba482808a8 100644 --- a/third_party/xla/xla/stream_executor/cuda/BUILD +++ b/third_party/xla/xla/stream_executor/cuda/BUILD @@ -84,21 +84,6 @@ config_setting( }, ) -bool_flag( - name = "enable_libnvjitlink_support", - build_setting_default = if_google( - True, - oss_value = False, - ), -) - -config_setting( - name = "libnvjitlink_support_enabled", - flag_values = { - ":enable_libnvjitlink_support": "True", - }, -) - cc_library( name = "cuda_platform_id", srcs = ["cuda_platform_id.cc"], @@ -1308,51 +1293,32 @@ cc_library( name = "nvjitlink_support", srcs = ["nvjitlink_support.cc"], hdrs = ["nvjitlink_support.h"], - local_defines = select({ - ":libnvjitlink_support_enabled": [ - "LIBNVJITLINK_SUPPORT=true", - ], - "//conditions:default": [ - "LIBNVJITLINK_SUPPORT=false", - ], - }) + if_cuda_newer_than( - "12_0", - ["CUDA_SUPPORTS_NVJITLINK=true"], - ["CUDA_SUPPORTS_NVJITLINK=false"], - ), ) cc_library( - name = "nvjitlink_stub", - srcs = [ + name = "nvjitlink", + srcs = ["nvjitlink.cc"], + hdrs = [ "nvjitlink.h", - "nvjitlink_stub.cc", ], - deps = [ - ":compilation_provider", - ":cuda_compute_capability", - "//xla/stream_executor/gpu:gpu_asm_opts", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", + local_defines = [ + # By default, nvJitLink.h defines static inline wrappers that forward to + # versioned __nvJitLink*__ symbols. Defining + # NVJITLINK_NO_INLINE makes nvjitlink.cc call the unversioned + # nvJitLink* symbols directly so it is compatible with nvjitlink_stub. + "NVJITLINK_NO_INLINE", ], -) - -cc_library( - name = "nvjitlink_impl", - srcs = [ - "nvjitlink.h", - "nvjitlink_impl.cc", + tags = [ + "cuda-only", + "gpu", ], - tags = ["manual"], deps = [ ":compilation_provider", ":cuda_compute_capability", ":ptx_compiler_helpers", "//xla/stream_executor:kernel_stats", "//xla/stream_executor/gpu:gpu_asm_opts", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", + "//xla/tsl/cuda:nvjitlink", # buildcleaner: keep "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/cleanup", @@ -1366,37 +1332,6 @@ cc_library( "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@local_config_cuda//cuda:cuda_headers", - "@local_config_cuda//cuda:nvjitlink", # buildcleaner: keep - ], -) - -# Since select() can't be nested, we need to wrap the cuda_newer_than check in a separate -# library target. -cc_library( - name = "nvjitlink_cuda_supported", - # Even though the macro is called `*_newer_than`, it does a greater-than-or-equal-to comparison. - deps = if_cuda_newer_than( - "12_0", - [":nvjitlink_impl"], - [":nvjitlink_stub"], - ), -) - -cc_library( - name = "nvjitlink", - hdrs = [ - "nvjitlink.h", - ], - deps = select({ - ":libnvjitlink_support_enabled": [":nvjitlink_cuda_supported"], - "//conditions:default": [":nvjitlink_stub"], - }) + [ - ":compilation_provider", - ":cuda_compute_capability", - "//xla/stream_executor:device_description", - "//xla/stream_executor/gpu:gpu_asm_opts", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/types:span", ], ) @@ -1412,19 +1347,20 @@ xla_cc_test( ]), # The test fails with msan/tsan since its introduction. tags = [ + "cuda-only", + "gpu", "nomsan", "notsan", ], deps = [ + ":compilation_provider", ":cuda_compute_capability", ":nvjitlink", ":nvjitlink_support", "//xla/stream_executor:kernel_stats", "//xla/stream_executor/gpu:gpu_asm_opts", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", - "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", @@ -1435,6 +1371,10 @@ cc_library( name = "nvjitlink_known_issues", srcs = ["nvjitlink_known_issues.cc"], hdrs = ["nvjitlink_known_issues.h"], + tags = [ + "cuda-only", + "gpu", + ], deps = [ ":nvjitlink", "@com_google_absl//absl/status:statusor", @@ -1446,15 +1386,17 @@ xla_cc_test( srcs = ["nvjitlink_known_issues_test.cc"], # LibNvJitLink is a binary-only library. Therefore is not compatible with msan/tsan. tags = [ + "cuda-only", + "gpu", "nomsan", "notsan", ], deps = [ + ":nvjitlink", ":nvjitlink_known_issues", ":nvjitlink_support", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", - "@tsl//tsl/platform:test", ], ) @@ -2448,6 +2390,10 @@ cc_library( name = "nvjitlink_compilation_provider", srcs = ["nvjitlink_compilation_provider.cc"], hdrs = ["nvjitlink_compilation_provider.h"], + tags = [ + "cuda-only", + "gpu", + ], deps = [ ":compilation_options", ":compilation_provider", @@ -2669,6 +2615,7 @@ cc_library( ":composite_compilation_provider", ":defer_relocatable_compilation_compilation_provider", ":driver_compilation_provider", + ":nvjitlink", ":nvjitlink_compilation_provider", ":nvjitlink_known_issues", ":nvjitlink_support", diff --git a/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc b/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc index 618c597b393d10..8b2956fe7e4381 100644 --- a/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc +++ b/third_party/xla/xla/stream_executor/cuda/assemble_compilation_provider.cc @@ -32,6 +32,7 @@ limitations under the License. #include "xla/stream_executor/cuda/composite_compilation_provider.h" #include "xla/stream_executor/cuda/defer_relocatable_compilation_compilation_provider.h" #include "xla/stream_executor/cuda/driver_compilation_provider.h" +#include "xla/stream_executor/cuda/nvjitlink.h" #include "xla/stream_executor/cuda/nvjitlink_compilation_provider.h" #include "xla/stream_executor/cuda/nvjitlink_known_issues.h" #include "xla/stream_executor/cuda/nvjitlink_support.h" @@ -64,6 +65,17 @@ absl::Status HasNvJitLinkSupport(const CompilationProviderOptions& options) { return absl::OkStatus(); } + // When built with the nvjitlink stub library, libnvJitLink.so is loaded + // dynamically at runtime and may not be present in the environment. Probe it + // via GetNvJitLinkVersion() first so auto mode can fall back gracefully. + if (absl::StatusOr version = GetNvJitLinkVersion(); + !version.ok()) { + return absl::UnavailableError( + absl::StrCat("LibNvJitLink is disabled in auto mode because the " + "library could not be loaded: ", + version.status().message())); + } + if (LoadedNvJitLinkHasKnownIssues()) { return absl::UnavailableError( "LibNvJitLink is disabled since the loaded library version has known " diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc similarity index 91% rename from third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc rename to third_party/xla/xla/stream_executor/cuda/nvjitlink.cc index 30056ad6b96853..9b043eab05fa1f 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "xla/stream_executor/cuda/nvjitlink.h" + #include #include #include @@ -39,12 +41,29 @@ limitations under the License. #include "third_party/gpus/cuda/include/nvJitLink.h" #include "xla/stream_executor/cuda/compilation_provider.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" -#include "xla/stream_executor/cuda/nvjitlink.h" #include "xla/stream_executor/cuda/ptx_compiler_helpers.h" #include "xla/stream_executor/gpu/gpu_asm_opts.h" #include "xla/stream_executor/kernel_stats.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" + +// Workaround for a bug in nvJitLink.h prior to CUDA 13.1: when +// NVJITLINK_NO_INLINE is defined, the header omits the `extern` declarations +// for the unversioned nvJitLink* functions instead of declaring them. This bug +// was fixed in CUDA 13.1. +// FIXME: Remove this workaround once the minimum supported CUDA version is at +// least 13.1. +extern "C" { +nvJitLinkResult nvJitLinkCreate(nvJitLinkHandle*, uint32_t, const char**); +nvJitLinkResult nvJitLinkDestroy(nvJitLinkHandle*); +nvJitLinkResult nvJitLinkAddData(nvJitLinkHandle, nvJitLinkInputType, + const void*, size_t, const char*); +nvJitLinkResult nvJitLinkComplete(nvJitLinkHandle); +nvJitLinkResult nvJitLinkGetLinkedCubinSize(nvJitLinkHandle, size_t*); +nvJitLinkResult nvJitLinkGetLinkedCubin(nvJitLinkHandle, void*); +nvJitLinkResult nvJitLinkGetErrorLogSize(nvJitLinkHandle, size_t*); +nvJitLinkResult nvJitLinkGetErrorLog(nvJitLinkHandle, char*); +nvJitLinkResult nvJitLinkGetInfoLogSize(nvJitLinkHandle, size_t*); +nvJitLinkResult nvJitLinkGetInfoLog(nvJitLinkHandle, char*); +} // extern "C" namespace stream_executor { diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc index 73b56b04fbe05e..03fe8a6e12f7ad 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink_known_issues_test.cc @@ -15,19 +15,24 @@ limitations under the License. #include "xla/stream_executor/cuda/nvjitlink_known_issues.h" +#include #include +#include "xla/stream_executor/cuda/nvjitlink.h" #include "xla/stream_executor/cuda/nvjitlink_support.h" -#include "tsl/platform/test.h" namespace { -TEST(NvJitLinkKnownIssuesTest, ReturnsFalseWhenNvJitLinkIsNotAvailable) { - if (stream_executor::IsLibNvJitLinkSupported()) { - GTEST_SKIP(); +TEST(NvJitLinkKnownIssuesTest, ReturnsExpectedResultForLoadedVersion) { + if (!stream_executor::IsLibNvJitLinkSupported()) { + EXPECT_FALSE(stream_executor::LoadedNvJitLinkHasKnownIssues()); + return; } - // This is the only invariance we can test without writing a pointless change - // detector test. - EXPECT_FALSE(stream_executor::LoadedNvJitLinkHasKnownIssues()); + ASSERT_OK_AND_ASSIGN(stream_executor::NvJitLinkVersion version, + stream_executor::GetNvJitLinkVersion()); + constexpr stream_executor::NvJitLinkVersion kMinVersionWithoutKnownIssues{12, + 5}; + EXPECT_EQ(stream_executor::LoadedNvJitLinkHasKnownIssues(), + version < kMinVersionWithoutKnownIssues); } } // namespace diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc deleted file mode 100644 index 4c3ddc4b5efd32..00000000000000 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_stub.cc +++ /dev/null @@ -1,40 +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 "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/types/span.h" -#include "xla/stream_executor/cuda/compilation_provider.h" -#include "xla/stream_executor/cuda/cuda_compute_capability.h" -#include "xla/stream_executor/cuda/nvjitlink.h" -#include "xla/stream_executor/gpu/gpu_asm_opts.h" - -namespace stream_executor { - -absl::StatusOr GetNvJitLinkVersion() { - return absl::UnimplementedError("libnvjitlink is not supported"); -} - -absl::StatusOr CompileAndLinkUsingLibNvJitLink( - const CudaComputeCapability& cc, absl::Span inputs, - GpuAsmOpts options, bool cancel_if_reg_spill, bool dump_compilation_log) { - return absl::UnimplementedError("libnvjitlink is not supported"); -} - -absl::StatusOr GetLatestPtxIsaVersionForLibNvJitLink() { - return absl::UnimplementedError("libnvjitlink is not supported"); -} - -} // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc index 0a1a870ff181da..48fbee2cf4cbf4 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink_support.cc @@ -24,7 +24,7 @@ bool IsLibNvJitLinkSupported() { #ifdef MEMORY_SANITIZER return false; #else - return LIBNVJITLINK_SUPPORT && CUDA_SUPPORTS_NVJITLINK; + return true; #endif } } // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc index 305b078601322d..4785528a129f78 100644 --- a/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc @@ -15,24 +15,20 @@ limitations under the License. #include "xla/stream_executor/cuda/nvjitlink.h" -#include - #include #include -#include #include #include #include #include "absl/status/status.h" #include "absl/status/status_matchers.h" // IWYU pragma: keep -#include "absl/strings/str_replace.h" #include "absl/types/span.h" +#include "xla/stream_executor/cuda/compilation_provider.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/cuda/nvjitlink_support.h" #include "xla/stream_executor/gpu/gpu_asm_opts.h" #include "xla/stream_executor/kernel_stats.h" -#include "xla/tsl/platform/statusor.h" namespace { @@ -43,7 +39,7 @@ namespace { // constexpr const char kDependeePtx[] = R"( .version 8.0 -.target sm_52 +.target sm_80 .address_size 64 // .globl _Z5magicv @@ -64,7 +60,7 @@ constexpr const char kDependeePtx[] = R"( // constexpr const char kDependentPtx[] = R"( .version 8.0 -.target sm_52 +.target sm_80 .address_size 64 // .globl _Z6kernelPi @@ -75,7 +71,6 @@ constexpr const char kDependentPtx[] = R"( .visible .entry _Z6kernelPi( .param .u64 _Z6kernelPi_param_0 ) -// Insert .maxnreg directive here! { .reg .b32 %r<2>; .reg .b64 %rd<3>; @@ -102,7 +97,7 @@ constexpr const char kDependentPtx[] = R"( // constexpr const char kStandalonePtx[] = R"( .version 8.0 -.target sm_52 +.target sm_80 .address_size 64 // .globl _Z6kernelPi @@ -123,8 +118,112 @@ constexpr const char kStandalonePtx[] = R"( })"; -constexpr stream_executor::CudaComputeCapability kDefaultComputeCapability{5, - 2}; +// Generated by the following command: +// +// echo "__global__ void kernel(int* in) { for (int i=0; i < 25; i++) \ +// { in[i] += i; } for (int i=0; i < 25; i++) { in[24-i] += i; }}" \ +// | nvcc -arch=sm_80 -o - -rdc true --ptx --x cu - +// +// The `.maxnreg` directive was added manually afterwards. +constexpr const char kSpillingPtx[] = R"( +.version 8.0 +.target sm_80 +.address_size 64 + + // .globl _Z6kernelPi + +.visible .entry _Z6kernelPi( + .param .u64 _Z6kernelPi_param_0 +) + .maxnreg 16 +{ + .reg .b32 %r<51>; + .reg .b64 %rd<3>; + + + ld.param.u64 %rd1, [_Z6kernelPi_param_0]; + cvta.to.global.u64 %rd2, %rd1; + ld.global.u32 %r1, [%rd2+4]; + ld.global.u32 %r2, [%rd2+8]; + ld.global.u32 %r3, [%rd2+12]; + ld.global.u32 %r4, [%rd2+16]; + ld.global.u32 %r5, [%rd2+20]; + ld.global.u32 %r6, [%rd2+24]; + ld.global.u32 %r7, [%rd2+28]; + ld.global.u32 %r8, [%rd2+32]; + ld.global.u32 %r9, [%rd2+36]; + ld.global.u32 %r10, [%rd2+40]; + ld.global.u32 %r11, [%rd2+44]; + ld.global.u32 %r12, [%rd2+48]; + ld.global.u32 %r13, [%rd2+52]; + ld.global.u32 %r14, [%rd2+56]; + ld.global.u32 %r15, [%rd2+60]; + ld.global.u32 %r16, [%rd2+64]; + ld.global.u32 %r17, [%rd2+68]; + ld.global.u32 %r18, [%rd2+72]; + ld.global.u32 %r19, [%rd2+76]; + ld.global.u32 %r20, [%rd2+80]; + ld.global.u32 %r21, [%rd2+84]; + ld.global.u32 %r22, [%rd2+88]; + ld.global.u32 %r23, [%rd2+92]; + ld.global.u32 %r24, [%rd2+96]; + add.s32 %r25, %r24, 24; + st.global.u32 [%rd2+96], %r25; + add.s32 %r26, %r23, 24; + st.global.u32 [%rd2+92], %r26; + add.s32 %r27, %r22, 24; + st.global.u32 [%rd2+88], %r27; + add.s32 %r28, %r21, 24; + st.global.u32 [%rd2+84], %r28; + add.s32 %r29, %r20, 24; + st.global.u32 [%rd2+80], %r29; + add.s32 %r30, %r19, 24; + st.global.u32 [%rd2+76], %r30; + add.s32 %r31, %r18, 24; + st.global.u32 [%rd2+72], %r31; + add.s32 %r32, %r17, 24; + st.global.u32 [%rd2+68], %r32; + add.s32 %r33, %r16, 24; + st.global.u32 [%rd2+64], %r33; + add.s32 %r34, %r15, 24; + st.global.u32 [%rd2+60], %r34; + add.s32 %r35, %r14, 24; + st.global.u32 [%rd2+56], %r35; + add.s32 %r36, %r13, 24; + st.global.u32 [%rd2+52], %r36; + add.s32 %r37, %r12, 24; + st.global.u32 [%rd2+48], %r37; + add.s32 %r38, %r11, 24; + st.global.u32 [%rd2+44], %r38; + add.s32 %r39, %r10, 24; + st.global.u32 [%rd2+40], %r39; + add.s32 %r40, %r9, 24; + st.global.u32 [%rd2+36], %r40; + add.s32 %r41, %r8, 24; + st.global.u32 [%rd2+32], %r41; + add.s32 %r42, %r7, 24; + st.global.u32 [%rd2+28], %r42; + add.s32 %r43, %r6, 24; + st.global.u32 [%rd2+24], %r43; + add.s32 %r44, %r5, 24; + st.global.u32 [%rd2+20], %r44; + add.s32 %r45, %r4, 24; + st.global.u32 [%rd2+16], %r45; + add.s32 %r46, %r3, 24; + st.global.u32 [%rd2+12], %r46; + add.s32 %r47, %r2, 24; + st.global.u32 [%rd2+8], %r47; + add.s32 %r48, %r1, 24; + st.global.u32 [%rd2+4], %r48; + ld.global.u32 %r49, [%rd2]; + add.s32 %r50, %r49, 24; + st.global.u32 [%rd2], %r50; + ret; +} +)"; + +constexpr stream_executor::CudaComputeCapability kDefaultComputeCapability{8, + 0}; // Just a helper function that wraps `CompileAndLinkUsingLibNvJitLink`. It makes // it easier to deal with C string PTX literals. @@ -187,13 +286,9 @@ TEST_F(NvJitLinkTest, CanAlsoJustCompileSingleCompilationUnit) { } TEST_F(NvJitLinkTest, CancelsOnRegSpill) { - std::string dependent_ptx = absl::StrReplaceAll( - kDependentPtx, {{"// Insert .maxnreg directive here!", ".maxnreg 16"}}); - // We have to disable optimization here, otherwise PTXAS will optimize our // trivial register usages away and we don't spill as intended. - EXPECT_THAT(CompileAndLinkHelper(kDefaultComputeCapability, - {dependent_ptx.c_str(), kDependeePtx}, + EXPECT_THAT(CompileAndLinkHelper(kDefaultComputeCapability, {kSpillingPtx}, /*disable_gpuasm_optimizations=*/true, /*cancel_if_reg_spill=*/true), absl_testing::StatusIs(absl::StatusCode::kCancelled)); @@ -201,8 +296,7 @@ TEST_F(NvJitLinkTest, CancelsOnRegSpill) { // We also test the converse to ensure our test case isn't broken. ASSERT_OK_AND_ASSIGN( stream_executor::cuda::Assembly assembly, - CompileAndLinkHelper(kDefaultComputeCapability, - {dependent_ptx.c_str(), kDependeePtx}, + CompileAndLinkHelper(kDefaultComputeCapability, {kSpillingPtx}, /*disable_gpuasm_optimizations=*/true, /*cancel_if_reg_spill=*/false)); ASSERT_EQ(assembly.module_stats.size(), 1); diff --git a/third_party/xla/xla/tsl/cuda/BUILD.bazel b/third_party/xla/xla/tsl/cuda/BUILD.bazel index 5b2105c7a56ba6..12b161e31ae14a 100644 --- a/third_party/xla/xla/tsl/cuda/BUILD.bazel +++ b/third_party/xla/xla/tsl/cuda/BUILD.bazel @@ -439,6 +439,55 @@ alias( visibility = ["//visibility:public"], ) +cuda_stub( + name = "nvjitlink", + srcs = ["nvjitlink.symbols"], +) + +cc_library( + name = "nvjitlink_stub", # buildifier: disable=duplicated-name + srcs = select({ + "@xla//xla/tsl:is_cuda_enabled_and_oss": [ + "nvjitlink.tramp.S", + "nvjitlink_stub.cc", + ], + "//conditions:default": [], + }), + defines = [ + # By default, nvJitLink.h defines static inline wrappers that forward to + # versioned __nvJitLink*__ symbols. Defining + # NVJITLINK_NO_INLINE makes consumers call the unversioned nvJitLink* + # symbols exported by nvjitlink.symbols instead. + "NVJITLINK_NO_INLINE", + ], + linkopts = if_cuda_is_configured(if_cuda_newer_than( + "13_0", + if_false = cuda_rpath_flags("nvidia/nvjitlink/lib"), + if_true = cuda_rpath_flags("nvidia/cu13/lib"), + )), + local_defines = [ + "IMPLIB_EXPORT_SHIMS=1", + ], + textual_hdrs = ["nvjitlink.inc"], + visibility = ["//visibility:public"], + deps = select({ + "@xla//xla/tsl:is_cuda_enabled_and_oss": [ + # keep sorted + "//xla/tsl/platform:logging", + "@local_config_cuda//cuda:cuda_headers", + "@tsl//tsl/platform:dso_loader", + "@tsl//tsl/platform:load_library", + ], + "//conditions:default": [], + }), +) + +alias( + name = "nvjitlink", # buildifier: disable=duplicated-name + actual = if_cuda_libs("@cuda_nvjitlink//:nvjitlink", ":nvjitlink_stub"), + visibility = ["//visibility:public"], +) + cuda_stub( name = "nvml", srcs = ["nvml.symbols"], diff --git a/third_party/xla/xla/tsl/cuda/nvjitlink.symbols b/third_party/xla/xla/tsl/cuda/nvjitlink.symbols new file mode 100644 index 00000000000000..805d74fb1f72c4 --- /dev/null +++ b/third_party/xla/xla/tsl/cuda/nvjitlink.symbols @@ -0,0 +1,16 @@ +nvJitLinkAddData +nvJitLinkAddFile +nvJitLinkComplete +nvJitLinkCreate +nvJitLinkDestroy +nvJitLinkGetErrorLog +nvJitLinkGetErrorLogSize +nvJitLinkGetInfoLog +nvJitLinkGetInfoLogSize +nvJitLinkGetLinkedCubin +nvJitLinkGetLinkedCubinSize +nvJitLinkGetLinkedLTOIR +nvJitLinkGetLinkedLTOIRSize +nvJitLinkGetLinkedPtx +nvJitLinkGetLinkedPtxSize +nvJitLinkVersion diff --git a/third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc b/third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc new file mode 100644 index 00000000000000..db753abb93e1c4 --- /dev/null +++ b/third_party/xla/xla/tsl/cuda/nvjitlink_stub.cc @@ -0,0 +1,73 @@ +/* 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 + +#include "third_party/gpus/cuda/include/nvJitLink.h" +#include "xla/tsl/platform/logging.h" +#include "tsl/platform/dso_loader.h" +#include "tsl/platform/load_library.h" + +// Implements the nvJitLink API by forwarding to nvJitLink loaded from the DSO. + +namespace { +// Returns DSO handle or null if loading the DSO fails. +void* GetDsoHandle() { + static auto handle = []() -> void* { + auto handle_or = tsl::internal::DsoLoader::GetNvJitLinkDsoHandle(); + if (!handle_or.ok()) { + return nullptr; + } + return handle_or.value(); + }(); + return handle; +} + +void* LoadSymbol(const char* symbol_name) { + void* symbol = nullptr; + if (auto handle = GetDsoHandle()) { + tsl::internal::GetSymbolFromLibrary(handle, symbol_name, &symbol) + .IgnoreError(); + } + return symbol; +} + +const char* kSymbols[] = { +#include "xla/tsl/cuda/nvjitlink.inc" +}; + +constexpr size_t kNumSymbols = sizeof(kSymbols) / sizeof(const char*); + +} // namespace + +extern "C" { + +static nvJitLinkResult GetSymbolNotFoundError() { + return NVJITLINK_ERROR_INTERNAL; +} + +extern void* _nvjitlink_tramp_table[]; + +void _nvjitlink_tramp_resolve(int i) { + CHECK_LE(0, i); + CHECK_LT(i, kNumSymbols); + void* p = LoadSymbol(kSymbols[i]); + if (!p) { + p = reinterpret_cast(&GetSymbolNotFoundError); + } + _nvjitlink_tramp_table[i] = p; +} + +} // extern "C" diff --git a/third_party/xla/xla/tsl/platform/default/dso_loader.cc b/third_party/xla/xla/tsl/platform/default/dso_loader.cc index ed9ecb415e5322..88e8c8afb36e5a 100644 --- a/third_party/xla/xla/tsl/platform/default/dso_loader.cc +++ b/third_party/xla/xla/tsl/platform/default/dso_loader.cc @@ -86,7 +86,8 @@ absl::StatusOr GetCudaDsoHandle(const std::string& name, if (name == "cufft" || name == "cusolver") { alt_version = (version == "12") ? "11" : (version == "11" ? "12" : ""); } else if (name == "cublas" || name == "cublasLt" || name == "cudart" || - name == "nvrtc" || name == "cusparse" || name == "cupti") { + name == "nvrtc" || name == "cusparse" || name == "cupti" || + name == "nvJitLink") { alt_version = (version == "13") ? "12" : (version == "12" ? "13" : ""); } @@ -128,6 +129,10 @@ absl::StatusOr GetNvrtcDsoHandle() { return GetCudaDsoHandle("nvrtc", GetCudaRtVersion()); } +absl::StatusOr GetNvJitLinkDsoHandle() { + return GetCudaDsoHandle("nvJitLink", GetCudaRtVersion()); +} + absl::StatusOr GetCudaRuntimeDsoHandle() { return GetCudaDsoHandle("cudart", GetCudaRtVersion()); } @@ -232,6 +237,11 @@ absl::StatusOr GetCudnnDsoHandle() { return *result; } +absl::StatusOr GetNvJitLinkDsoHandle() { + static auto result = new auto(DsoLoader::GetNvJitLinkDsoHandle()); + return *result; +} + } // namespace CachedDsoLoader } // namespace internal } // namespace tsl diff --git a/third_party/xla/xla/tsl/platform/default/dso_loader.h b/third_party/xla/xla/tsl/platform/default/dso_loader.h index 32e80ef473ec78..ffbfcb48d876d8 100644 --- a/third_party/xla/xla/tsl/platform/default/dso_loader.h +++ b/third_party/xla/xla/tsl/platform/default/dso_loader.h @@ -16,6 +16,10 @@ limitations under the License. // Common DSO loading functionality: exposes callables that dlopen DSOs // in either the runfiles directories +// IWYU pragma: private, include "tsl/platform/dso_loader.h" +// IWYU pragma: friend ".*tsl/platform/.*dso_loader\.h" +// IWYU pragma: friend ".*tsl/platform/default/.*" + #ifndef XLA_TSL_PLATFORM_DEFAULT_DSO_LOADER_H_ #define XLA_TSL_PLATFORM_DEFAULT_DSO_LOADER_H_ @@ -41,6 +45,7 @@ absl::StatusOr GetNcclDsoHandle(); absl::StatusOr GetNvshmemDsoHandle(); absl::StatusOr GetNvInferDsoHandle(); absl::StatusOr GetNvInferPluginDsoHandle(); +absl::StatusOr GetNvJitLinkDsoHandle(); absl::StatusOr GetNvmlDsoHandle(); absl::StatusOr GetNvrtcDsoHandle(); @@ -69,6 +74,7 @@ absl::StatusOr GetCusolverDsoHandle(); absl::StatusOr GetCusparseDsoHandle(); absl::StatusOr GetCuptiDsoHandle(); absl::StatusOr GetCudnnDsoHandle(); +absl::StatusOr GetNvJitLinkDsoHandle(); } // namespace CachedDsoLoader } // namespace internal From c0a46ed3fd48dc99b25d23b3140192031f23f4ee Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Tue, 29 Sep 2026 05:35:22 -0700 Subject: [PATCH 10/25] Enable strict properties in assembly format for tf_framework dialect MLIR is migrating dialects to strict property handling in custom assembly formats (`useStrictPropertiesInAssemblyFormat = 1`), where `attr-dict` only parses and prints discardable attributes and inherent attributes must be either explicitly bound in the assembly format or handled via `prop-dict`. Add `prop-dict` before `attr-dict` in the assembly formats of `tf_framework.alloc` and `tf_framework.jit_compile_from_str`, enable `useStrictPropertiesInAssemblyFormat = 1` on `TFFramework_Dialect`, and update tests to use the `<...>` property dictionary syntax. PiperOrigin-RevId: 990243949 --- .../tools/kernel_gen/ir/tf_framework_ops.td | 6 +-- .../kernel_gen/tests/embed_tf_framework.mlir | 8 ++-- .../tests/func_to_jit_invocations.mlir | 46 +++++++++---------- .../mlir/tools/kernel_gen/tests/ops.mlir | 22 ++++----- .../tests/tf_framework_legalize_to_llvm.mlir | 4 +- 5 files changed, 43 insertions(+), 43 deletions(-) diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td b/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td index b5e11c53ddd97e..0f54049616b5ad 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td +++ b/tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td @@ -40,7 +40,7 @@ def TFFramework_Dialect : Dialect { }]; let useDefaultTypePrinterParser = 1; - let useStrictPropertiesInAssemblyFormat = 0; + let useStrictPropertiesInAssemblyFormat = 1; } def TFFramework_OpKernelContextType : DialectType !tf_framework.jit_callable attributes {tf_entry} { // CHECK: %[[RES:.*]] = tf_framework.jit_compile_from_str %[[CTX]], "placeholder" // CHECK: return %[[RES]] - %0 = tf_framework.jit_compile_from_str "placeholder" { - architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3], - unrollFactors = [4], enableFtz = false, - index64Bit = false, cpuCodegen = false } + %0 = tf_framework.jit_compile_from_str "placeholder" < + tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false, + index64Bit = false, cpuCodegen = false> { + architectures = ["sm_123", "sm_456"] } func.return %0 : !tf_framework.jit_callable } diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir b/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir index d19e8180973b76..b9abece1fee264 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir @@ -41,12 +41,12 @@ func.func @unary_tanh(%arg : tensor) -> tensor { // CHECK-SAME: } // CHECK-SAME: } // CHECK-SAME: " -// CHECK-SAME: { -// CHECK-SAME: cpuCodegen = false -// CHECK-SAME: enableFtz = false +// CHECK-SAME: < // CHECK-SAME: tileSizes = [1, 2, 3] // CHECK-SAME: unrollFactors = [3, 2, 1] -// CHECK-SAME: } +// CHECK-SAME: enableFtz = false +// CHECK-SAME: cpuCodegen = false +// CHECK-SAME: > // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]]) // CHECK: return %[[RES]] @@ -59,11 +59,11 @@ func.func @unary_tanh(%arg : tensor) -> tensor { // CHECK-JFLT: %[[IF:.*]] = scf.if %[[CMPI]] // CHECK-JFLT: %[[JIT:.*]] = tf_framework.jit_compile_from_str // CHECK-JFLT-SAME: "module -// CHECK-JFLT-SAME: cpuCodegen = false -// CHECK-JFLT-SAME: enableFtz = false -// CHECK-JFLT-SAME: index64Bit = true // CHECK-JFLT-SAME: tileSizes = [1, 2, 3] // CHECK-JFLT-SAME: unrollFactors = [3, 2, 1] +// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: index64Bit = true +// CHECK-JFLT-SAME: cpuCodegen = false // CHECK-JFLT: %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]]) // CHECK-JFLT: scf.yield %[[JIT_0]] // CHECK-JFLT: else @@ -91,12 +91,12 @@ func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*x // CHECK-SAME: } // CHECK-SAME: } // CHECK-SAME: " -// CHECK-SAME: { -// CHECK-SAME: cpuCodegen = false -// CHECK-SAME: enableFtz = false +// CHECK-SAME: < // CHECK-SAME: tileSizes = [1, 2, 3] // CHECK-SAME: unrollFactors = [3, 2, 1] -// CHECK-SAME: } +// CHECK-SAME: enableFtz = false +// CHECK-SAME: cpuCodegen = false +// CHECK-SAME: > // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]]) // CHECK: return %[[RES]] @@ -122,12 +122,12 @@ func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*x // CHECK-JFLT-SAME: } // CHECK-JFLT-SAME: } // CHECK-JFLT-SAME: " -// CHECK-JFLT-SAME: { -// CHECK-JFLT-SAME: cpuCodegen = false -// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: < // CHECK-JFLT-SAME: tileSizes = [1, 2, 3] // CHECK-JFLT-SAME: unrollFactors = [3, 2, 1] -// CHECK-JFLT-SAME: } +// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: cpuCodegen = false +// CHECK-JFLT-SAME: > // CHECK-JFLT: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]]) // CHECK-JFLT: scf.yield %[[RES]] : tensor<*xf32> // CHECK-JFLT: } else { @@ -155,13 +155,13 @@ func.func @reciprocal(%arg0: tensor<*xf32>) // CHECK=SAME: return %[[RES_JIT]] : tensor<*xf32> // CHECK-SAME: } // CHECK-SAME: } -// CHECK-SAME: { -// CHECK-SAME: cpuCodegen = false, +// CHECK-SAME: < +// CHECK-SAME: tileSizes = [1, 2, 3], +// CHECK-SAME: unrollFactors = [3, 2, 1], // CHECK-SAME: enableFtz = false, // CHECK-SAME: index64Bit = false, -// CHECK-SAME: tileSizes = [1, 2, 3], -// CHECK-SAME: unrollFactors = [3, 2, 1] -// CHECK-SAME: } +// CHECK-SAME: cpuCodegen = false +// CHECK-SAME: > // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]]) : tensor<*xf32> -> tensor<*xf32> // CHECK: return %[[RES]] : tensor<*xf32> @@ -174,11 +174,11 @@ func.func @reciprocal(%arg0: tensor<*xf32>) // CHECK-JFLT: %[[IF:.*]] = scf.if %[[CMPI]] // CHECK-JFLT: %[[JIT:.*]] = tf_framework.jit_compile_from_str // CHECK-JFLT-SAME: "module -// CHECK-JFLT-SAME: cpuCodegen = false -// CHECK-JFLT-SAME: enableFtz = false -// CHECK-JFLT-SAME: index64Bit = true // CHECK-JFLT-SAME: tileSizes = [1, 2, 3] // CHECK-JFLT-SAME: unrollFactors = [3, 2, 1] +// CHECK-JFLT-SAME: enableFtz = false +// CHECK-JFLT-SAME: index64Bit = true +// CHECK-JFLT-SAME: cpuCodegen = false // CHECK-JFLT: %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]]) // CHECK-JFLT: scf.yield %[[JIT_0]] // CHECK-JFLT: else diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir b/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir index 8d988e2eae11bc..5829355053fa35 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir @@ -29,10 +29,10 @@ func.func @alloc(%ctx: !tf_framework.op_kernel_context, // CHECK-LABEL: func @forwarding_alloc func.func @forwarding_alloc(%ctx: !tf_framework.op_kernel_context, %size_0 : index , %size_2 : index) { - %buf = tf_framework.alloc(%ctx, %size_0, %size_2) { + %buf = tf_framework.alloc(%ctx, %size_0, %size_2) < input_indices = [0 : i32, 1 : i32], - output_index = 0 : i32 - } : memref + output_index = 0 + > : memref func.return } @@ -94,20 +94,20 @@ func.func @jit_compile(%ctx : !tf_framework.op_kernel_context) // CHECK-LABEL: func @jit_compile_from_str_wo_ctx func.func @jit_compile_from_str_wo_ctx() -> !tf_framework.jit_callable { - %callable = tf_framework.jit_compile_from_str "placeholder" { - architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3], - unrollFactors = [4], enableFtz = false, - index64Bit = false, cpuCodegen = false } + %callable = tf_framework.jit_compile_from_str "placeholder" < + tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false, + index64Bit = false, cpuCodegen = false> { + architectures = ["sm_123", "sm_456"] } func.return %callable : !tf_framework.jit_callable } // CHECK-LABEL: func @jit_compile_from_str func.func @jit_compile_from_str(%ctx : !tf_framework.op_kernel_context) -> !tf_framework.jit_callable { - %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" { - architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3], - unrollFactors = [4], enableFtz = false, - index64Bit = false, cpuCodegen = false } + %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" < + tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false, + index64Bit = false, cpuCodegen = false> { + architectures = ["sm_123", "sm_456"] } func.return %callable : !tf_framework.jit_callable } diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir b/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir index 15e0284d89eb20..e6794dcf88936a 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir @@ -206,9 +206,9 @@ func.func @jit_compile_from_str(%ctx: !tf_framework.op_kernel_context) // CHECK-SAME: %[[NUM_UNROLL_FACTORS]], %[[UNROLL_FACTORS]], // CHECK-SAME: %[[ENABLE_FTZ]], %[[CPU_CODEGEN]] // CHECK: llvm.return %[[RES]] - %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" { + %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" < tileSizes = [1, 2, 3], unrollFactors = [4], - enableFtz = false, index64Bit = false, cpuCodegen = false } + enableFtz = false, index64Bit = false, cpuCodegen = false> func.return %0 : !tf_framework.jit_callable } From 014be1c569f75890ece90914b90c4cc74bb815ea Mon Sep 17 00:00:00 2001 From: Ce Zheng CA Date: Tue, 29 Sep 2026 05:37:13 -0700 Subject: [PATCH 11/25] [XLA] Add 0-arg root_piece() overloads for LiteralBase data and size accessors `LiteralBase::data`, `LiteralBase::untyped_data`, `LiteralBase::size_bytes`, `LiteralBase::total_size_bytes`, and `LiteralBase::element_count` previously took a default argument `const ShapeIndex& shape_index = {}`, causing every 0-arg caller across XLA to construct and destroy a temporary `ShapeIndex` (`absl::InlinedVector`) at the call site and route through `piece(shape_index)` tree traversal. Similarly, `MutableLiteralBase::PopulateParallel` and `PopulateLinearParallel` checked `data().size() > 32` at every template instantiation site (instantiating `data()` just to obtain the element count) rather than `element_count() > 32` like `PopulateInplaceParallel` and `PopulateLinearInplaceParallel`. Split the 0-arg overloads of `data()`, `untyped_data()`, `size_bytes()`, `total_size_bytes()`, and `element_count()` to access `root_piece()` / `mutable_root_piece()` directly without constructing a temporary `ShapeIndex`, define `element_count()` and `element_count(const ShapeIndex&)` out-of-line in `literal.cc` delegating to `root_piece().element_count()` / `piece(index).element_count()`, and use `element_count() > 32` in `PopulateParallel` and `PopulateLinearParallel`. ### Runtime & Memory Performance Impact - **CPU / Latency**: Positive. 0-arg calls to `data()`, `untyped_data()`, `size_bytes()`, `total_size_bytes()`, and `element_count()` bypass temporary `ShapeIndex` (`absl::InlinedVector`) construction/destruction and `piece(shape_index)` loop traversal, accessing `root_piece()` directly in O(1) with zero heap or stack-vector overhead. - **Memory**: Neutral to positive (eliminates temporary `ShapeIndex` stack frames at call sites). PiperOrigin-RevId: 990244740 --- third_party/xla/xla/literal.cc | 24 ++++++++++++++++++ third_party/xla/xla/literal.h | 46 +++++++++++++++++++++------------- 2 files changed, 53 insertions(+), 17 deletions(-) diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc index 160ad73cd42b76..bc09effad37ee4 100644 --- a/third_party/xla/xla/literal.cc +++ b/third_party/xla/xla/literal.cc @@ -2957,22 +2957,46 @@ LiteralProto LiteralBase::ToProto() const { return proto; } +const void* LiteralBase::untyped_data() const { + return root_piece().untyped_data(); +} + const void* LiteralBase::untyped_data(const ShapeIndex& shape_index) const { return piece(shape_index).untyped_data(); } +void* MutableLiteralBase::untyped_data() { + return mutable_root_piece().untyped_data(); +} + void* MutableLiteralBase::untyped_data(const ShapeIndex& shape_index) { return piece(shape_index).untyped_data(); } +int64_t LiteralBase::size_bytes() const { + return root_piece().size_bytes_dense(); +} + int64_t LiteralBase::size_bytes(const ShapeIndex& shape_index) const { return piece(shape_index).size_bytes_dense(); } +int64_t LiteralBase::total_size_bytes() const { + return root_piece().total_bytes_dense(); +} + int64_t LiteralBase::total_size_bytes(const ShapeIndex& shape_index) const { return piece(shape_index).total_bytes_dense(); } +int64_t LiteralBase::element_count() const { + return root_piece().element_count(); +} + +int64_t LiteralBase::element_count(const ShapeIndex& index) const { + return piece(index).element_count(); +} + std::string LiteralBase::GetR1U8AsString() const { CHECK(shape().IsArray()); CHECK_EQ(shape().dimensions().size(), 1); diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h index 9a90f698d96d76..5c9d91bd2c608c 100644 --- a/third_party/xla/xla/literal.h +++ b/third_party/xla/xla/literal.h @@ -104,16 +104,21 @@ class LiteralBase { // ShapeIndex is not array. See primitive_util.h for the mapping from XLA type // to native type. template - absl::Span data(const ShapeIndex& shape_index = {}) const; + absl::Span data() const; + template + absl::Span data(const ShapeIndex& shape_index) const; // Returns a const pointer to (or size of) the underlying buffer holding the // array at the given shape index. CHECKs if the subshape of the literal at // the given ShapeIndex is not array. - const void* untyped_data(const ShapeIndex& shape_index = {}) const; - int64_t size_bytes(const ShapeIndex& shape_index = {}) const; + const void* untyped_data() const; + const void* untyped_data(const ShapeIndex& shape_index) const; + int64_t size_bytes() const; + int64_t size_bytes(const ShapeIndex& shape_index) const; // total size in bytes of the literal (including pre-allocated dynamic // metadata) - int64_t total_size_bytes(const ShapeIndex& shape_index = {}) const; + int64_t total_size_bytes() const; + int64_t total_size_bytes(const ShapeIndex& shape_index) const; // Computes the size in bytes of the output of the Serialize method. // If `pack_pred` is true, PRED elements are bit-packed (8 elements per byte); @@ -380,13 +385,8 @@ class LiteralBase { // Returns the count of the elements in the array at the given shape index in // this literal. - int64_t element_count(const ShapeIndex& index = {}) const { - if (index.empty()) { - // Common case, avoid GetSubshape(). - return ShapeUtil::ElementsIn(shape()); - } - return ShapeUtil::ElementsIn(ShapeUtil::GetSubshape(shape(), index)); - } + int64_t element_count() const; + int64_t element_count(const ShapeIndex& index) const; // This definition is here to ensure that nobody accidentally implements this // function which would lead to inconsistencies. Use Hash instead. @@ -1309,7 +1309,9 @@ class MutableLiteralBase : public LiteralBase { // given ShapeIndex is not array. See primitive_util.h for the mapping from // XLA type to native type. template - absl::Span data(const ShapeIndex& shape_index = {}); + absl::Span data(); + template + absl::Span data(const ShapeIndex& shape_index); // Unhide const method from parent class. using LiteralBase::data; @@ -1325,7 +1327,8 @@ class MutableLiteralBase : public LiteralBase { // Returns a pointer to the underlying buffer holding the array at the given // shape index. CHECKs if the subshape of the literal at the given ShapeIndex // is not array. - void* untyped_data(const ShapeIndex& shape_index = {}); + void* untyped_data(); + void* untyped_data(const ShapeIndex& shape_index); // Unhide const method from parent class. using LiteralBase::untyped_data; @@ -1944,12 +1947,22 @@ void LiteralBase::Piece::SetLinear(int64_t linear_index, NativeT value) { data().data()[linear_index] = value; } +template +absl::Span LiteralBase::data() const { + return root_piece().data(); +} + template absl::Span LiteralBase::data( const ShapeIndex& shape_index) const { return piece(shape_index).data(); } +template +absl::Span MutableLiteralBase::data() { + return mutable_root_piece().data(); +} + template absl::Span MutableLiteralBase::data(const ShapeIndex& shape_index) { return piece(shape_index).data(); @@ -2208,7 +2221,7 @@ absl::Status MutableLiteralBase::PopulateParallel(Generator&& generator) { TF_RET_CHECK(shape().IsArray()) << __func__ << " is only supported for dense arrays: " << shape(); return PopulateInternal(generator, - /*parallel=*/data().size() > 32); + /*parallel=*/element_count() > 32); } template @@ -2247,9 +2260,8 @@ template ( - std::forward(generator), - /*parallel=*/data().size() > 32); + return PopulateLinearInternal(std::forward(generator), + /*parallel=*/element_count() > 32); } template *> From 2a994e6353dbeeab9a1f7f3f5d385d1151842cfe Mon Sep 17 00:00:00 2001 From: Mohammed Anany Date: Tue, 29 Sep 2026 06:19:41 -0700 Subject: [PATCH 12/25] [XLA:GPU] Evaluate tiling candidates in parallel in SoftmaxRewriterTriton PiperOrigin-RevId: 990263277 --- .../xla/xla/backends/gpu/transforms/BUILD | 3 ++ .../gpu/transforms/softmax_rewriter_triton.cc | 19 +++++--- .../gpu/transforms/softmax_rewriter_triton.h | 15 +++++- .../softmax_rewriter_triton_test.cc | 46 +++++++++++++++++++ .../xla/xla/service/gpu/gpu_compiler.cc | 3 +- 5 files changed, 77 insertions(+), 9 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 4c4089e2c614d0..09a84ce25588e7 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -3280,6 +3280,7 @@ cc_library( "//xla/service/gpu/model:triton_emitter_constraints", "//xla/stream_executor:device_description", "//xla/tools:hlo_decomposer_lib", + "//xla/tsl/platform:env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -3311,8 +3312,10 @@ xla_cc_test( "//xla/service/gpu:alias_info", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_device_info_for_tests", + "//xla/service/gpu/model:gpu_indexing_performance_model", "//xla/stream_executor:device_description", "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tsl/platform:env", "//xla/tsl/platform:errors", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc index cbc5735abd8af8..d16ae5e238a76d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc @@ -66,6 +66,7 @@ limitations under the License. #include "xla/shape.h" #include "xla/stream_executor/device_description.h" #include "xla/tools/hlo_decomposer.h" +#include "xla/tsl/platform/threadpool.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -353,12 +354,15 @@ DecideIfShouldFuseAndMaybeSetBlockLevelParameters( const se::DeviceDescription& device_info, const HloCostAnalysis::ShapeSizeFunction& shape_size, const GpuAliasInfo* alias_info, MLIRContext* mlir_context, - bool use_cost_model_to_evaluate_fusions) { + bool use_cost_model_to_evaluate_fusions, + tsl::thread::ThreadPool* thread_pool = nullptr) { auto fusion_adaptor = HloFusionAdaptor::ForInstruction(normalization_fusion); ABSL_ASSIGN_OR_RETURN(TiledRunTimeDataOrError tiled_runtime_data_or, indexing_performance_model - .TryFindBestTilingForFusionAsync(*fusion_adaptor) + .TryFindBestTilingForFusionAsync( + *fusion_adaptor, + thread_pool ? thread_pool->AsExecutor() : nullptr) .Await()); if (const auto* fusion_decision = @@ -407,7 +411,8 @@ absl::StatusOr MaybeFuseDiamondImpl( const se::DeviceDescription& device_info, const HloCostAnalysis::ShapeSizeFunction& shape_size, const GpuAliasInfo* alias_info, MLIRContext* mlir_context, - bool use_cost_model_to_evaluate_fusions) { + bool use_cost_model_to_evaluate_fusions, + tsl::thread::ThreadPool* thread_pool = nullptr) { ABSL_ASSIGN_OR_RETURN(HloFusionInstruction * normalization_fusion, MakeFusionForDiamond(diamond)); HloInstruction* root = diamond.root; @@ -418,7 +423,7 @@ absl::StatusOr MaybeFuseDiamondImpl( DecideIfShouldFuseAndMaybeSetBlockLevelParameters( normalization_fusion, indexing_performance_model, device_info, shape_size, alias_info, mlir_context, - use_cost_model_to_evaluate_fusions)); + use_cost_model_to_evaluate_fusions, thread_pool)); if (fusion_decision.IsForbidden()) { VLOG(2) << "Not fusing: " << fusion_decision.Explain(); @@ -675,11 +680,13 @@ absl::StatusOr SoftmaxRewriterTriton::MaybeFuseNormalizationDiamond( diamond.root->GetModule() ->config() .debug_options() - .xla_gpu_experimental_enable_same_shape_multi_output_fusion()); + .xla_gpu_experimental_enable_same_shape_multi_output_fusion(), + mlir_context_pool_); return MaybeFuseDiamondImpl(diamond, indexing_performance_model, device_info_, shape_size_, alias_info_, mlir_context_, - use_cost_model_to_evaluate_fusions_); + use_cost_model_to_evaluate_fusions_, + thread_pool_); } absl::StatusOr SoftmaxRewriterTriton::RunImpl( diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h index 426dc73c7ef912..19b442ab94efa9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h +++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.h @@ -27,10 +27,15 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" #include "xla/service/gpu/alias_info.h" +#include "xla/service/gpu/model/gpu_indexing_performance_model.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/stream_executor/device_description.h" +namespace tsl::thread { +class ThreadPool; +} // namespace tsl::thread + namespace xla { namespace gpu { @@ -55,13 +60,17 @@ class SoftmaxRewriterTriton : public HloModulePass { const GpuAliasInfo* alias_info, mlir::MLIRContext* mlir_context, bool only_fuse_if_profitable, - bool use_experimental_tiling) + bool use_experimental_tiling, + tsl::thread::ThreadPool* thread_pool = nullptr, + MlirContextPool* mlir_context_pool = nullptr) : device_info_(device_info), shape_size_(shape_size), alias_info_(alias_info), use_cost_model_to_evaluate_fusions_(only_fuse_if_profitable), use_experimental_tiling_(use_experimental_tiling), - mlir_context_(mlir_context) {} + mlir_context_(mlir_context), + thread_pool_(thread_pool), + mlir_context_pool_(mlir_context_pool) {} absl::string_view name() const override { return "triton-softmax-rewriter"; } @@ -114,6 +123,8 @@ class SoftmaxRewriterTriton : public HloModulePass { bool use_cost_model_to_evaluate_fusions_; bool use_experimental_tiling_; mlir::MLIRContext* mlir_context_; + tsl::thread::ThreadPool* thread_pool_ = nullptr; + MlirContextPool* mlir_context_pool_ = nullptr; }; } // namespace gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc index 6af936e36b83f7..182225da938efb 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 @@ -35,12 +35,15 @@ limitations under the License. #include "xla/service/gpu/alias_info.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/service/gpu/model/gpu_indexing_performance_model.h" #include "xla/service/hlo_cost_analysis.h" #include "xla/service/instruction_fusion.h" #include "xla/service/pattern_matcher.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/threadpool.h" #include "xla/xla.pb.h" namespace xla { @@ -1165,6 +1168,49 @@ ENTRY main { .WithPredicate(HasBlockLevelFusionConfig))); } +TEST_P(SoftmaxRewriterTritonTest, CanFuseWithParallelTilingSearch) { + const std::string hlo_string = R"( +HloModule softmax +max_computation { + arg_0 = f32[] parameter(0) + arg_1 = f32[] parameter(1) + ROOT maximum = f32[] maximum(arg_0, arg_1) +} +add_computation { + arg_0 = f32[] parameter(0) + arg_1 = f32[] parameter(1) + ROOT add = f32[] add(arg_0, arg_1) +} +ENTRY main { + param_0 = f32[127,125]{1,0} parameter(0) + constant_neg_inf = f32[] constant(-inf) + reduce = f32[127]{0} reduce(param_0, constant_neg_inf), dimensions={1}, to_apply=max_computation + broadcast = f32[127,125]{1,0} broadcast(reduce), dimensions={0} + ROOT subtract = f32[127,125]{1,0} subtract(param_0, broadcast) +})"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), "test_pool", 4); + // Mirrors the contexts GpuCompiler pools: multithreading is disabled, so the + // cost model must give each candidate its own context. + MlirContextPool mlir_context_pool( + [] { + return std::make_unique( + mlir::MLIRContext::Threading::DISABLED); + }, + /*preallocate=*/4); + SoftmaxRewriterTriton rewriter( + device_info_, HloCostAnalysis::DefaultShapeSize, &alias_info_, + &mlir_context_, + /*only_fuse_if_profitable=*/false, + /*use_experimental_tiling=*/GetParam(), &thread_pool, &mlir_context_pool); + EXPECT_THAT(rewriter.Run(module.get()), IsOkAndHolds(true)); + EXPECT_TRUE(verifier().Run(module.get()).status().ok()); + EXPECT_THAT( + module->entry_computation()->root_instruction(), + GmockMatch( + m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig))); +} + INSTANTIATE_TEST_SUITE_P(SoftmaxRewriterTritonTestSuite, SoftmaxRewriterTritonTest, ::testing::Bool(), [](const ::testing::TestParamInfo& info) { diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index b2c99e37d950ae..6a828612bf9bb7 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -2193,7 +2193,8 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment( alias_info, mlir_context, /*only_fuse_if_profitable=*/true, /*use_experimental_tiling=*/ - debug_options.xla_gpu_experimental_enable_tiling_propagation()); + debug_options.xla_gpu_experimental_enable_tiling_propagation(), + thread_pool, &mlir_context_pool_); } pipeline.AddPass(); From d8812599b68be2260a77fc8f723eb5c3c6bfb555 Mon Sep 17 00:00:00 2001 From: Ce Zheng CA Date: Tue, 29 Sep 2026 06:32:13 -0700 Subject: [PATCH 13/25] [XLA] Move 2-arg IndexUtil::MultidimensionalIndexToLinearIndex out-of-line to index_util.cc The 2-argument convenience overload `IndexUtil::MultidimensionalIndexToLinearIndex(const Shape&, absl::Span)` was defined inline in `index_util.h`, causing `LayoutUtil::MinorToMajor(shape)`, `Shape::array_state_maybe_underneath_buffer()` checks, and the stride accumulation loop to be inlined into hundreds of `LiteralBase::Piece::Get` and `Set` template instantiations. Meanwhile, hot-path inner loops in `LiteralBase::Broadcast` and `HloEvaluator` use the 3-argument fast-path overload `IndexUtil::MultidimensionalIndexToLinearIndex(const Shape&, absl::Span minor_to_major, absl::Span multi_index)` where `minor_to_major` is pre-extracted outside the loop. Move the 2-argument overload out-of-line to `index_util.cc`, keep the 3-argument fast-path overload inline in `index_util.h`, and hoist `absl::Span dims = shape.dimensions();` outside the stride accumulation loop in the 3-argument overload so `shape.dimensions(...)` is not re-evaluated on every dimension. ### Runtime & Memory Performance Impact - **CPU / Latency**: Preserves hot-path inlining of the 3-argument overload in `LiteralBase::Broadcast` and `HloEvaluator` (with a slight speedup on higher-rank broadcasts from hoisting `shape.dimensions()` out of the inner loop): ``` Benchmark Base CPU (ns) New CPU (ns) Delta -------------------------------------------------------------------------------- BM_BroadcastVectorToMatrix/16/16 3950 4086 -3.4% BM_BroadcastVectorToMatrix/16/1024 133732 131341 +1.8% BM_BroadcastVectorToMatrix/1024/1024 8646962 8415307 +2.7% BM_BroadcastMatrixToTensor/16/16 19883 18694 +6.0% BM_BroadcastMatrixToTensor/16/1024 1145118 1109215 +3.1% BM_BroadcastMatrixToTensor/1024/1024 73361100 68770331 +6.3% BM_Populate/64/process_time 52482 61484 -17.1% BM_Populate/1024/process_time 13224583 14506707 -9.7% BM_PopulateParallel/1024/process_time 23279761 26892908 -15.5% BM_PopulateLinear/1024/process_time 2123731 2081302 +2.0% BM_PopulateLinearParallel/1024/process_time 3329581 3376813 -1.4% ``` - **Memory**: Neutral (no changes to object layouts or heap allocations). PiperOrigin-RevId: 990268395 --- third_party/xla/xla/index_util.cc | 7 +++++++ third_party/xla/xla/index_util.h | 10 ++++------ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/index_util.cc b/third_party/xla/xla/index_util.cc index 7377db56f9d814..d2a018ad003a1f 100644 --- a/third_party/xla/xla/index_util.cc +++ b/third_party/xla/xla/index_util.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/log/check.h" #include "absl/types/span.h" #include "xla/layout_util.h" #include "xla/shape.h" @@ -26,6 +27,12 @@ limitations under the License. namespace xla { +/* static */ int64_t IndexUtil::MultidimensionalIndexToLinearIndex( + const Shape& shape, absl::Span multi_index) { + return MultidimensionalIndexToLinearIndex( + shape, LayoutUtil::MinorToMajor(shape), multi_index); +} + /* static */ DimensionVector IndexUtil::LinearIndexToMultidimensionalIndex( const Shape& shape, int64_t linear_index) { DCHECK_GE(linear_index, 0); diff --git a/third_party/xla/xla/index_util.h b/third_party/xla/xla/index_util.h index ca8ce5e924d1da..9cf51389414087 100644 --- a/third_party/xla/xla/index_util.h +++ b/third_party/xla/xla/index_util.h @@ -42,10 +42,7 @@ class IndexUtil { // on the shape and its layout. The first index in the multi_index is // dimension 0. static int64_t MultidimensionalIndexToLinearIndex( - const Shape& shape, absl::Span multi_index) { - return MultidimensionalIndexToLinearIndex( - shape, LayoutUtil::MinorToMajor(shape), multi_index); - } + const Shape& shape, absl::Span multi_index); // Converts a multidimensional index (eg {x, y, z}) into a linear index based // on the shape and its layout. The first index in the multi_index is @@ -94,9 +91,10 @@ class IndexUtil { // D{L(1)} * ... // Scale factor holding the growing product of D{L(i)} terms. + absl::Span dims = shape.dimensions(); for (size_t i = 0; i < multi_index.size(); ++i) { DCHECK_GE(multi_index[i], 0); - DCHECK_LT(multi_index[i], shape.dimensions(i)) + DCHECK_LT(multi_index[i], dims[i]) << "indexing beyond extent in dimension " << i << ":" << "\n\tindex: " << absl::StrJoin(multi_index, ",") << "\n\tshape: " << ShapeUtil::HumanString(shape); @@ -107,7 +105,7 @@ class IndexUtil { int64_t linear_index = multi_index[minor_to_major[0]]; int64_t scale = 1; for (int i = 1; i < minor_to_major.size(); ++i) { - scale *= shape.dimensions(minor_to_major[i - 1]); + scale *= dims[minor_to_major[i - 1]]; linear_index += scale * multi_index[minor_to_major[i]]; } return linear_index; From c532f8fb7a0e7cee6aad11a87e70113e0466f23d Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Tue, 29 Sep 2026 06:36:05 -0700 Subject: [PATCH 14/25] Update `xla/AGENTS.md` with GPU testing and license header guidelines Add guidance to `xla/AGENTS.md` covering: - Passing `--config=cuda` when running GPU tests. - Including the OpenXLA Apache 2.0 copyright header with the current calendar year on newly created source and `BUILD` files. PiperOrigin-RevId: 990270007 --- third_party/xla/xla/AGENTS.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/third_party/xla/xla/AGENTS.md b/third_party/xla/xla/AGENTS.md index 1d945983879635..58e5d6c4077e71 100644 --- a/third_party/xla/xla/AGENTS.md +++ b/third_party/xla/xla/AGENTS.md @@ -92,6 +92,7 @@ suggesting, or modifying code within the OpenXLA codebase (everything under * Use `HloPjRtInterpreterReferenceMixin` or `HloHardwareIndependentTestBase` for compiler pass tests locally where possible. + * When running GPU tests, pass `--config=cuda` to the test command. * Ensure tests are deterministic and do not flake. 6. **BUILD targets**: @@ -117,6 +118,30 @@ suggesting, or modifying code within the OpenXLA codebase (everything under * **Always** use the static `OpTy::create(rewriter, ...)` method when creating MLIR operations. * **Avoid** using `rewriter.create(...)`. This syntax is deprecated. +11. **License Headers**: + * Any new source or build files created in this codebase must include the + OpenXLA Apache 2.0 copyright header at the top of the file. + * Replace `` in the template below with the current 4-digit calendar + year (e.g., `2026`—do not leave `` literally or copy an outdated + year): + ```cpp + /* Copyright 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. + ==============================================================================*/ + ``` + (Use `#` comment prefixes for Python and `BUILD` files.) + ## Editing This File This file is shared between the open-source repository and Google-internal From bc2ea0ed374949482f210e87cdd26f4d2707d903 Mon Sep 17 00:00:00 2001 From: "Patrick C. Toulme" Date: Tue, 29 Sep 2026 06:58:15 -0700 Subject: [PATCH 15/25] [XLA] Cut copying and repeated visits in HloReplicationAnalysis The two trivial replication states own no heap object. Replication trees hold no Shape, propagation merges in place, and post orders are computed once. A computation already marked unique is not marked again, and repeat visits evaluate only instructions whose inputs changed. In a computation marked unique, tuple roots and SPMD infeeds are now unique too; other results are unchanged. PiperOrigin-RevId: 990279838 --- third_party/xla/xla/hlo/analysis/BUILD | 2 + .../hlo/analysis/hlo_replication_analysis.cc | 367 +++++++++++---- .../hlo/analysis/hlo_replication_analysis.h | 126 ++++- .../analysis/hlo_replication_analysis_test.cc | 444 ++++++++++++++++++ 4 files changed, 823 insertions(+), 116 deletions(-) diff --git a/third_party/xla/xla/hlo/analysis/BUILD b/third_party/xla/xla/hlo/analysis/BUILD index d0b2e163a1da12..952d475c49271a 100644 --- a/third_party/xla/xla/hlo/analysis/BUILD +++ b/third_party/xla/xla/hlo/analysis/BUILD @@ -381,12 +381,14 @@ cc_library( deps = [ "//xla:shape_tree", "//xla:shape_util", + "//xla:tuple_tree", "//xla:util", "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "@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/container:node_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc index 5443f988a4991e..6045dec29aa0ce 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc @@ -46,6 +46,7 @@ limitations under the License. #include "xla/map_util.h" #include "xla/shape_tree.h" #include "xla/shape_util.h" +#include "xla/tuple_tree.h" #include "xla/xla_data.pb.h" namespace xla { @@ -310,70 +311,212 @@ HloReplicationAnalysis::DetermineHloInstructionIsReplicated( return HloReplication::UniqueOnAllDevices(); } +const TupleTree& +HloReplicationAnalysis::GetReplication(const HloInstruction* inst) const { + auto it = hlo_replication_.find(inst); + CHECK(it != hlo_replication_.end()) + << inst->name() << " has no replication yet"; + return it->second; +} + +bool HloReplicationAnalysis::CombineReplication( + const TupleTree& source, TupleTree* dest) { + bool updated = false; + // Both trees are built from compatible shapes, so their pre order nodes line + // up one to one. + auto source_node = source.begin(); + for (TupleTree::NodePair& node : *dest) { + CHECK(source_node != source.end()) + << "source tree ends before " << node.first; + CHECK(source_node->first == node.first) + << source_node->first << " vs " << node.first; + HloReplication new_replication = + MergeReplications(node.second, source_node->second); + if (!node.second.Equal(new_replication)) { + node.second = std::move(new_replication); + updated = true; + } + ++source_node; + } + return updated; +} + +bool HloReplicationAnalysis::AssignOrCombineReplication( + TupleTree replication, const HloInstruction* dest) { + auto [it, inserted] = + hlo_replication_.try_emplace(dest, std::move(replication)); + // try_emplace leaves `replication` untouched when the key exists. + if (inserted || CombineReplication(replication, &it->second)) { + OnReplicationChanged(dest); + return true; + } + return false; +} + +bool HloReplicationAnalysis::PropagateReplication(const HloInstruction* source, + const HloInstruction* dest) { + auto source_it = hlo_replication_.find(source); + if (source_it == hlo_replication_.end()) { + return false; + } + auto dest_it = hlo_replication_.find(dest); + if (dest_it != hlo_replication_.end()) { + if (!CombineReplication(source_it->second, &dest_it->second)) { + return false; + } + } else { + // Copy before inserting: the insertion may rehash the map. + TupleTree copy(source_it->second); + hlo_replication_.try_emplace(dest, std::move(copy)); + } + OnReplicationChanged(dest); + return true; +} + +bool HloReplicationAnalysis::MarkNotReplicated(const HloInstruction* inst) { + auto [it, inserted] = hlo_replication_.try_emplace( + inst, inst->shape(), HloReplication::UniqueOnAllDevices()); + if (inserted) { + OnReplicationChanged(inst); + return true; + } + // Merging any value with UniqueOnAllDevices() gives UniqueOnAllDevices(), so + // the existing tree is overwritten in place; only nodes that were not unique + // yet count as a change for the fixed point. + bool updated = false; + for (TupleTree::NodePair& node : it->second) { + if (!node.second.IsUniqueOnAllDevices()) { + node.second = HloReplication::UniqueOnAllDevices(); + updated = true; + } + } + if (updated) { + OnReplicationChanged(inst); + } + return updated; +} + +void HloReplicationAnalysis::OnReplicationChanged(const HloInstruction* inst) { + const HloComputation* computation = inst->parent(); + // With partial replication every visit evaluates every instruction. Users + // come later in post order, so a visit that evaluates every instruction of + // the computation reaches them, and only a change of its root matters. + if (support_partial_replication_ || + (computation == visiting_computation_ && visiting_all_instructions_ && + inst != computation->root_instruction())) { + return; + } + MarkDependentsDirty(inst); +} + +void HloReplicationAnalysis::MarkDependentsDirty(const HloInstruction* inst) { + const HloComputation* computation = inst->parent(); + if (inst == computation->root_instruction()) { + for (const HloInstruction* caller : computation->caller_instructions()) { + MarkDirty(caller); + } + } + // The first visit of a computation evaluates every instruction too. + if (computation == visiting_computation_ && visiting_all_instructions_) { + return; + } + auto it = computation_states_.find(computation); + if (it == computation_states_.end() || !it->second.visited) { + return; + } + for (const HloInstruction* user : inst->users()) { + MarkDirty(user); + } +} + +void HloReplicationAnalysis::MarkDirty(const HloInstruction* inst) { + // An instruction without a replication has not been evaluated yet. A visit + // that evaluates every instruction reaches it later, or is evaluating it now + // and MarkEvaluated then checks the computations it calls. + if (!hlo_replication_.contains(inst) || !dirty_.insert(inst).second) { + return; + } + const HloComputation* computation = inst->parent(); + if (computation_states_[computation].num_dirty++ == 0) { + for (const HloInstruction* caller : computation->caller_instructions()) { + MarkDirty(caller); + } + } +} + +void HloReplicationAnalysis::MarkEvaluated(const HloInstruction* inst) { + for (const HloComputation* callee : inst->called_computations()) { + auto it = computation_states_.find(callee); + if (it != computation_states_.end() && it->second.num_dirty > 0) { + MarkDirty(inst); + return; + } + } + if (dirty_.erase(inst) > 0) { + --computation_states_[inst->parent()].num_dirty; + } +} + bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( const HloComputation* computation, bool mark_everything_not_replicated) { + // A repeat visit with everything marked unique cannot change anything. + if (mark_everything_not_replicated && + !computations_marked_not_replicated_.insert(computation).second) { + return false; + } + ComputationState& state = computation_states_[computation]; + if (state.post_order.empty()) { + state.post_order = computation->MakeInstructionPostOrder(); + } + // A repeat visit that does not mark evaluates only the dirty instructions: + // merging a value that is replicated or unique on all devices with itself + // changes nothing. A partially replicated value whose device sets are all + // singletons merges with itself to unique, so with partial replication every + // visit evaluates every instruction. + const bool evaluate_all = !state.visited || mark_everything_not_replicated || + support_partial_replication_; + if (!evaluate_all && state.num_dirty == 0) { + return false; + } + const HloComputation* outer_computation = visiting_computation_; + const bool outer_evaluates_all = visiting_all_instructions_; + visiting_computation_ = computation; + visiting_all_instructions_ = evaluate_all; bool changed = false; - for (const HloInstruction* inst : computation->MakeInstructionPostOrder()) { - // Assigns the shape tree to dest if dest doesn't have one yet, or combines - // it with the existing one by and'ing them. Returns if anything is updated. - auto assign_or_combine_shapetree = - [&](ShapeTree&& to_combine, - const HloInstruction* dest) { - auto it = hlo_replication_.find(dest); - if (it == hlo_replication_.end()) { - hlo_replication_[dest] = std::move(to_combine); - return true; - } - bool updated = false; - it->second.ForEachMutableElement( - [&](const ShapeIndex& index, HloReplication* element) { - HloReplication new_replication = - MergeReplications(*element, to_combine.element(index)); - if (!element->Equal(new_replication)) { - *element = std::move(new_replication); - updated = true; - } - }); - return updated; - }; - // Assigns or combines source's shape tree to dest. Returns if anything is - // updated. - auto propagate_shapetree = [&](const HloInstruction* source, - const HloInstruction* dest) { - auto source_it = hlo_replication_.find(source); - if (source_it == hlo_replication_.end()) { - return false; - } - return assign_or_combine_shapetree( - ShapeTree(source_it->second), dest); - }; + for (const HloInstruction* inst : state.post_order) { + if (!evaluate_all && !dirty_.contains(inst)) { + continue; + } + // Whether the evaluation visits the computations that `inst` calls. + bool visits_computations = false; // For the opcodes below that we do special handling, we don't need to // explicitly check mark_everything_not_replicated because if it is set, the // operands should already be marked as not replicated. if (inst->opcode() == HloOpcode::kWhile) { + visits_computations = true; // Since while body's input and output alias each other, we need to run it // multiple times until a fixed point is reached. while (true) { // First, propagate the input's and body root's shape trees to the // parameters of the body and condition. - bool updated = propagate_shapetree( + bool updated = PropagateReplication( inst->operand(0), inst->while_condition()->parameter_instruction(0)); - updated |= propagate_shapetree( + updated |= PropagateReplication( inst->while_body()->root_instruction(), inst->while_condition()->parameter_instruction(0)); - updated |= propagate_shapetree( + updated |= PropagateReplication( inst->operand(0), inst->while_body()->parameter_instruction(0)); updated |= - propagate_shapetree(inst->while_body()->root_instruction(), - inst->while_body()->parameter_instruction(0)); + PropagateReplication(inst->while_body()->root_instruction(), + inst->while_body()->parameter_instruction(0)); // Compute the condition. updated |= ComputeHloReplicationOnComputation( inst->while_condition(), mark_everything_not_replicated); // Compute the body. If the condition is not replicated, the while body // should be different across replicas. if (!ContainsKey(loops_known_with_same_iterations_, inst) && - !hlo_replication_[inst->while_condition()->root_instruction()] + !GetReplication(inst->while_condition()->root_instruction()) .element({}) .IsReplicatedOnAllDevices()) { updated |= ComputeHloReplicationOnComputation( @@ -388,11 +531,12 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( changed = true; } // Propagate the input's and body root's shape trees to the while HLO. - changed |= propagate_shapetree(inst->operand(0), inst); + changed |= PropagateReplication(inst->operand(0), inst); changed |= - propagate_shapetree(inst->while_body()->root_instruction(), inst); + PropagateReplication(inst->while_body()->root_instruction(), inst); } else if (inst->opcode() == HloOpcode::kCall || inst->opcode() == HloOpcode::kFusion) { + visits_computations = true; auto called = inst->called_computations().front(); // If the called computation has several call sites, conservatively mark // everything inside as not-replicated, since we can't reason about @@ -403,23 +547,24 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( /*mark_everything_not_replicated=*/true); } else { for (int64_t i = 0; i < inst->operand_count(); ++i) { - changed |= propagate_shapetree(inst->operand(i), - called->parameter_instruction(i)); + changed |= PropagateReplication(inst->operand(i), + called->parameter_instruction(i)); } changed |= ComputeHloReplicationOnComputation( called, mark_everything_not_replicated); } - changed |= propagate_shapetree(called->root_instruction(), inst); + changed |= PropagateReplication(called->root_instruction(), inst); } else if (inst->opcode() == HloOpcode::kConditional) { + visits_computations = true; // Propagate inputs' shape trees to the called computations' parameters. for (int64_t i = 0; i < inst->called_computations().size(); ++i) { - changed |= propagate_shapetree( + changed |= PropagateReplication( inst->operand(i + 1), inst->called_computations()[i]->parameter_instruction(0)); } // If the condition is not replicated, the conditional result should be // different across replicas. - if (!hlo_replication_[inst->operand(0)] + if (!GetReplication(inst->operand(0)) .element({}) .IsReplicatedOnAllDevices()) { for (auto called : inst->called_computations()) { @@ -427,64 +572,67 @@ bool HloReplicationAnalysis::ComputeHloReplicationOnComputation( called, /*mark_everything_not_replicated=*/true); } - changed |= assign_or_combine_shapetree( - ShapeTree(inst->shape(), - HloReplication::UniqueOnAllDevices()), - inst); + changed |= MarkNotReplicated(inst); } else { for (auto called : inst->called_computations()) { changed |= ComputeHloReplicationOnComputation( called, mark_everything_not_replicated); - changed |= propagate_shapetree(called->root_instruction(), inst); + changed |= PropagateReplication(called->root_instruction(), inst); } } } else if (inst->opcode() == HloOpcode::kTuple) { - ShapeTree shape_tree( - inst->shape(), HloReplication::ReplicatedOnAllDevices()); + // When marking, the tuple is unique at its own index like every other + // value of the computation; its elements come from marked operands. + TupleTree tree( + inst->shape(), mark_everything_not_replicated + ? HloReplication::UniqueOnAllDevices() + : HloReplication::ReplicatedOnAllDevices()); for (int64_t i = 0; i < inst->operand_count(); ++i) { - shape_tree.CopySubtreeFrom(hlo_replication_[inst->operand(i)], {}, {i}); + CHECK_OK(tree.CopyCompatibleSubtreeFrom( + GetReplication(inst->operand(i)), {}, {i})); } - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + changed |= AssignOrCombineReplication(std::move(tree), inst); } else if (inst->opcode() == HloOpcode::kOptimizationBarrier) { - ShapeTree shape_tree = hlo_replication_[inst->operand(0)]; - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + changed |= PropagateReplication(inst->operand(0), inst); } else if (inst->opcode() == HloOpcode::kGetTupleElement) { - ShapeTree shape_tree( - inst->shape(), HloReplication::ReplicatedOnAllDevices()); - shape_tree.CopySubtreeFrom(hlo_replication_[inst->operand(0)], - {inst->tuple_index()}, {}); - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); - } else if (inst->opcode() == HloOpcode::kInfeed && cross_partition_spmd_) { - ShapeTree shape_tree( - inst->shape(), HloReplication::UniqueOnAllDevices()); + TupleTree tree(inst->shape(), + HloReplication::ReplicatedOnAllDevices()); + CHECK_OK(tree.CopyCompatibleSubtreeFrom(GetReplication(inst->operand(0)), + {inst->tuple_index()}, {})); + changed |= AssignOrCombineReplication(std::move(tree), inst); + } else if (inst->opcode() == HloOpcode::kInfeed && cross_partition_spmd_ && + !mark_everything_not_replicated) { + // Replicated elements are equal on every partition only if all of them + // run the infeed equally often, so marked computations skip this rule. + TupleTree tree(inst->shape(), + HloReplication::UniqueOnAllDevices()); if (inst->has_sharding()) { auto sharding = inst->sharding().GetAsShapeTree(inst->shape()); - shape_tree.ForEachMutableElement( + tree.ForEachMutableElement( [&sharding](const ShapeIndex& index, HloReplication* data) { *data = sharding.element(index).IsReplicated() ? HloReplication::ReplicatedOnAllDevices() : HloReplication::UniqueOnAllDevices(); }); } - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + changed |= AssignOrCombineReplication(std::move(tree), inst); + } else if (mark_everything_not_replicated) { + changed |= MarkNotReplicated(inst); } else { - if (mark_everything_not_replicated) { - changed |= assign_or_combine_shapetree( - ShapeTree(inst->shape(), - HloReplication::UniqueOnAllDevices()), - inst); - } else { - ShapeTree shape_tree( - inst->shape(), HloReplication::ReplicatedOnAllDevices()); - ShapeUtil::ForEachSubshape( - inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { - *shape_tree.mutable_element(index) = - DetermineHloInstructionIsReplicated(inst, index); - }); - changed |= assign_or_combine_shapetree(std::move(shape_tree), inst); + TupleTree tree(inst->shape(), + HloReplication::ReplicatedOnAllDevices()); + for (TupleTree::NodePair& node : tree) { + node.second = DetermineHloInstructionIsReplicated(inst, node.first); } + changed |= AssignOrCombineReplication(std::move(tree), inst); + } + if (state.num_dirty > 0 || visits_computations) { + MarkEvaluated(inst); } } + state.visited = true; + visiting_computation_ = outer_computation; + visiting_all_instructions_ = outer_evaluates_all; return changed; } @@ -495,8 +643,8 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { auto entry = module_->entry_computation(); for (int i = 0; i < entry->num_parameters(); ++i) { auto param = entry->parameter_instruction(i); - ShapeTree shape_tree(param->shape(), - HloReplication::UniqueOnAllDevices()); + TupleTree tree(param->shape(), + HloReplication::UniqueOnAllDevices()); std::unique_ptr> sharding_tree = nullptr; if (cross_partition_spmd_) { @@ -524,7 +672,7 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { if (sharding_tree != nullptr) { // In cross-partition spmd mode, set parameter replication status // based on the parameter's sharding. - *shape_tree.mutable_element(index) = + *tree.mutable_element(index) = sharding_tree->element(index).IsReplicated() ? HloReplication::ReplicatedOnAllDevices() : HloReplication::UniqueOnAllDevices(); @@ -535,13 +683,13 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { if (!cross_partition_spmd_ && (*replication)[leaf_index]) { // Setting parameter replication status for replicas in // non cross-partition spmd mode. - *shape_tree.mutable_element(index) = + *tree.mutable_element(index) = HloReplication::ReplicatedOnAllDevices(); } if (cross_partition_spmd_ && !(*replication)[leaf_index]) { // Setting paramemter replication status for partitions in // cross-partition spmd mode. - *shape_tree.mutable_element(index) = + *tree.mutable_element(index) = HloReplication::UniqueOnAllDevices(); } ++leaf_index; @@ -549,10 +697,14 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { return absl::OkStatus(); }); ABSL_RETURN_IF_ERROR(status); - hlo_replication_[param] = std::move(shape_tree); + hlo_replication_[param] = std::move(tree); } ComputeHloReplicationOnComputation(entry, /*mark_everything_not_replicated=*/false); + // Scratch state of the walk. + computation_states_.clear(); + dirty_.clear(); + computations_marked_not_replicated_.clear(); return absl::OkStatus(); } @@ -679,25 +831,27 @@ HloReplicationAnalysis::HloReplication::HloReplication() : state_(State::kReplicatedOnAllDevices) {} HloReplicationAnalysis::HloReplication::HloReplication( - HloReplicationAnalysis::HloReplication::State state, + HloReplicationAnalysis::HloReplication::State state) + : state_(state) { + DCHECK(state != State::kPartiallyReplicated); +} + +HloReplicationAnalysis::HloReplication::HloReplication( absl::Span> device_set_root_per_replica) - : state_(state), + : state_(State::kPartiallyReplicated), device_set_root_per_replica_( std::make_shared< HashOnConstruction>>>( - device_set_root_per_replica)) { - CHECK(state == State::kPartiallyReplicated || - device_set_root_per_replica_->empty()); -} + device_set_root_per_replica)) {} HloReplicationAnalysis::HloReplication HloReplicationAnalysis::HloReplication::ReplicatedOnAllDevices() { - return HloReplication(State::kReplicatedOnAllDevices, {}); + return HloReplication(State::kReplicatedOnAllDevices); } HloReplicationAnalysis::HloReplication HloReplicationAnalysis::HloReplication::UniqueOnAllDevices() { - return HloReplication(State::kUniqueOnAllDevices, {}); + return HloReplication(State::kUniqueOnAllDevices); } HloReplicationAnalysis::HloReplication @@ -724,8 +878,7 @@ HloReplicationAnalysis::HloReplication::PartiallyReplicated( } device_set_root_per_replica.push_back(std::move(device_set_root)); } - return HloReplication(State::kPartiallyReplicated, - device_set_root_per_replica); + return HloReplication(device_set_root_per_replica); } HloReplicationAnalysis::HloReplication @@ -781,9 +934,14 @@ HloReplicationAnalysis::HloReplication::Merge( bool HloReplicationAnalysis::HloReplication::Equal( const HloReplication& other) const { - return state_ == other.state_ && - device_set_root_per_replica_->hash_ == - other.device_set_root_per_replica_->hash_; + if (state_ != other.state_) { + return false; + } + if (state_ != State::kPartiallyReplicated) { + return true; + } + return device_set_root_per_replica_->hash_ == + other.device_set_root_per_replica_->hash_; } bool HloReplicationAnalysis::HloReplication::operator==( @@ -799,9 +957,16 @@ bool HloReplicationAnalysis::HloReplication::IsUniqueOnAllDevices() const { return state_ == State::kUniqueOnAllDevices; } +bool HloReplicationAnalysis::HloReplication::IsPartiallyReplicated() const { + return state_ == State::kPartiallyReplicated; +} + bool HloReplicationAnalysis::HloReplication::IsReplicatedWithinSubgroup( absl::Span device_ids) const { if (device_ids.empty()) return true; + if (state_ != State::kPartiallyReplicated) { + return state_ == State::kReplicatedOnAllDevices; + } for (const std::vector& device_set_roots : *device_set_root_per_replica_) { if (!absl::c_all_of(device_ids, diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h index d050d44e19c9f5..3c4f8eb6844ae4 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.h @@ -26,11 +26,15 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/container/node_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/shape_util.h" +#include "xla/tuple_tree.h" #include "xla/xla_data.pb.h" namespace xla { @@ -123,13 +127,17 @@ class HloReplicationAnalysis { bool operator==(const HloReplication& rhs) const; bool IsReplicatedOnAllDevices() const; bool IsUniqueOnAllDevices() const; + bool IsPartiallyReplicated() const; bool IsReplicatedWithinSubgroup(absl::Span device_ids) const; std::string ToString() const; template friend H AbslHashValue(H h, const HloReplication& r) { - return H::combine(std::move(h), r.state_, - *r.device_set_root_per_replica_); + h = H::combine(std::move(h), r.state_); + if (r.device_set_root_per_replica_ != nullptr) { + h = H::combine(std::move(h), *r.device_set_root_per_replica_); + } + return h; } private: @@ -138,8 +146,10 @@ class HloReplicationAnalysis { kUniqueOnAllDevices = 1, kPartiallyReplicated = 2, }; + // Only partially replicated values carry device sets; the other two states + // share a null pointer so that copying them is a plain state copy. + explicit HloReplication(State state); explicit HloReplication( - State state, absl::Span> device_set_root_per_replica); State state_; // Helper class that subclasses T, and computes the hash once on @@ -161,15 +171,17 @@ class HloReplicationAnalysis { return H::combine(std::move(h), r.hash_); } }; - // Empty if state_ is kReplicatedOnAllDevices or kUniqueOnAllDevices. - - // If cross_partition_spmd is true, groups_for_replicas_[k]'s size equals - // the number of partitions, and within replica k, groups_for_replicas_[k] - // maps each partition ID to the smallest partition ID in the set. + // Null if state_ is kReplicatedOnAllDevices or kUniqueOnAllDevices. // - // If cross_partition_spmd is false, groups_for_replicas_[k]'s size equals - // the number of replicas, and within partition k, groups_for_replicas_[k] - // maps each replica to the smallest replica ID in the set. + // If cross_partition_spmd is true, device_set_root_per_replica_[k]'s size + // equals the number of partitions, and within replica k, + // device_set_root_per_replica_[k] maps each partition ID to the smallest + // partition ID in the set. + // + // If cross_partition_spmd is false, device_set_root_per_replica_[k]'s size + // equals the number of replicas, and within partition k, + // device_set_root_per_replica_[k] maps each replica to the smallest replica + // ID in the set. std::shared_ptr>>> device_set_root_per_replica_; }; @@ -182,6 +194,13 @@ class HloReplicationAnalysis { HloReplication MergeReplications(const HloReplication& replication_a, const HloReplication& replication_b) { + // Merging with a value that is replicated or unique on all devices is a + // copy of one side; only merges of two partially replicated values are + // worth memoizing. + if (!replication_a.IsPartiallyReplicated() || + !replication_b.IsPartiallyReplicated()) { + return replication_a.Merge(replication_b); + } std::pair key = {replication_a, replication_b}; @@ -214,6 +233,47 @@ class HloReplicationAnalysis { bool ComputeHloReplicationOnComputation(const HloComputation* computation, bool mark_everything_not_replicated); + // Records that the replication of `inst` changed: its users, and the callers + // of its computation if it is the root, have to be evaluated again. + void OnReplicationChanged(const HloInstruction* inst); + + // Marks the users of `inst`, and the callers of its computation if it is the + // root, dirty, unless the visit in progress evaluates them anyway. + void MarkDependentsDirty(const HloInstruction* inst); + + // Adds `inst` to dirty_ if it has been evaluated before. The first dirty + // instruction of a computation also makes the callers of that computation + // dirty, since the pending evaluation runs when a caller visits the + // computation again. + void MarkDirty(const HloInstruction* inst); + + // Called after `inst` was evaluated: `inst` stays or becomes dirty if a + // computation it calls has dirty instructions, and leaves dirty_ otherwise. + void MarkEvaluated(const HloInstruction* inst); + + // Returns the replication of `inst`, which must have been computed already. + const TupleTree& GetReplication( + const HloInstruction* inst) const; + + // Merges `source` into `dest` element by element. Returns whether anything + // changed. + bool CombineReplication(const TupleTree& source, + TupleTree* dest); + + // Assigns `replication` to `dest` if it has none yet, or combines it with + // the existing one. Returns whether anything changed. + bool AssignOrCombineReplication(TupleTree replication, + const HloInstruction* dest); + + // Assigns or combines the replication of `source` to `dest`. Returns whether + // anything changed; nothing changes if `source` has no replication yet. + bool PropagateReplication(const HloInstruction* source, + const HloInstruction* dest); + + // Marks `inst` unique on all devices at every index. Returns whether + // anything changed. + bool MarkNotReplicated(const HloInstruction* inst); + // Builds the replica group dedup map that allows caching replication // calculations for all-reduce/all-gather that share the same replica groups. // This can significantly help in compile times when replica groups are very @@ -243,12 +303,48 @@ class HloReplicationAnalysis { // Capture the number of partitions / replicas for the module. const int64_t num_partitions_, replica_count_; - // A map from each analyzed HLO instruction to a shape tree that represents - // whether the instruction outputs the same value across replicas or - // partitions at each shape index. - absl::flat_hash_map> + // A map from each analyzed HLO instruction to a tree that represents whether + // the instruction outputs the same value across replicas or partitions at + // each shape index. The trees hold no Shape, so they stay valid when callers + // change shapes or delete instructions while they hold the analysis. + absl::flat_hash_map> hlo_replication_; + struct ComputationState { + // The module does not change during the analysis, so the post order is + // computed once. + std::vector post_order; + // Whether a visit of the computation has completed. + bool visited = false; + // Number of instructions of the computation in dirty_. + int64_t num_dirty = 0; + }; + // The values must not move when the map grows, since a visit of a nested + // computation adds an entry while the caller iterates over its post order. + absl::node_hash_map + computation_states_; + + // Instructions whose operands, or the roots of the computations they call, + // changed after their last evaluation, and instructions that call a + // computation with dirty instructions. Without partial replication, a repeat + // visit of a computation that is not marked evaluates only these: any other + // instruction would compute its old value again. + absl::flat_hash_set dirty_; + + // The computation of the innermost visit in progress, and whether that visit + // evaluates every instruction. + const HloComputation* visiting_computation_ = nullptr; + bool visiting_all_instructions_ = false; + + // Computations already visited with mark_everything_not_replicated set. + // That visit marks the parameters and every ordinary instruction unique, + // infeeds included; tuples, get-tuple-elements, optimization barriers and + // nested calls derive from those values, and replication only moves towards + // unique, so a repeat visit with the flag set would recompute the same values + // and is skipped. + absl::flat_hash_set + computations_marked_not_replicated_; + // Replications for all-reduce/all-gather that have the same replica groups is // usually identical. We use the following data structures to memoize the // replications for instructions with identical replica groups. diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc index 6c601626141e72..e9d6575e3545b6 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/hlo/analysis/hlo_replication_analysis.h" +#include #include #include #include @@ -1078,5 +1079,448 @@ ENTRY entry { FindInstruction(module.get(), "gte.1"), {})); } +// Loop carried values that are partially replicated with two different +// groupings, unique, and replicated, at their own tuple indices. The loop +// needs a second iteration to settle: the partially replicated all-reduce +// results only reach the parameter after the first pass over the body. +TEST_F(HloReplicationAnalysisTest, WhileLoopPartialReplicationPerIndex) { + const std::string module_str = R"hlo( +HloModule WhileLoopPartialReplicationPerIndex + +sum { + a = f32[] parameter(0) + b = f32[] parameter(1) + ROOT add = f32[] add(a, b) +} + +cond { + cond_param = (f32[8], f32[8], f32[8], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=3 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[8], f32[8], f32[8], u32[]) parameter(0) + x = f32[8] get-tuple-element(body_param), index=0 + y = f32[8] get-tuple-element(body_param), index=1 + u = f32[8] get-tuple-element(body_param), index=2 + i = u32[] get-tuple-element(body_param), index=3 + xu = f32[8] add(x, u) + yu = f32[8] add(y, u) + ar0 = f32[8] all-reduce(xu), to_apply=sum, replica_groups={{0,1},{2,3}} + ar1 = f32[8] all-reduce(yu), to_apply=sum, replica_groups={{0,2},{1,3}} + mixed = f32[8] add(ar0, ar1) + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[8], f32[8], f32[8], u32[]) tuple(ar0, ar1, mixed, next_i) +} + +ENTRY entry { + p0 = f32[8] parameter(0), parameter_replication={true} + p1 = f32[8] parameter(1), parameter_replication={true} + p2 = f32[8] parameter(2), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[8], f32[8], f32[8], u32[]) tuple(p0, p1, p2, zero) + ROOT while = (f32[8], f32[8], f32[8], u32[]) while(init), condition=cond, body=body +} +)hlo"; + const std::vector pairs01_23 = + CreateReplicaGroups({{0, 1}, {2, 3}}); + const std::vector pairs02_13 = + CreateReplicaGroups({{0, 2}, {1, 3}}); + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/4)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::RunWithPartialReplication( + module.get(), /*cross_partition_spmd=*/false)); + + for (const char* name : {"while", "body_param", "tuple"}) { + const HloInstruction* inst = FindInstruction(module.get(), name); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {0})) << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {0}, pairs01_23)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {0}, pairs02_13)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {1})) << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {1}, pairs01_23)) + << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {1}, pairs02_13)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {2})) << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {2}, pairs01_23)) + << name; + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(inst, {2}, pairs02_13)) + << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {3})) << name; + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt(inst, {3}, pairs01_23)) + << name; + } + // u is unique, so the all-reduce operands are unique in every pass and each + // all-reduce result is replicated within its own groups only. + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "xu"), {}, pairs01_23)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "yu"), {}, pairs02_13)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "ar0"), {}, pairs01_23)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "ar1"), {}, pairs02_13)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "mixed"), {}, pairs01_23)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "mixed"), {}, pairs02_13)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "lt"), {})); +} + +// The loop state enters through a tuple instruction, whose tree is replicated +// at the tuple root as well as at the leaves. The condition turns +// non replicated after the first pass, and the body parameter must then be +// unique at every index, the tuple root included. +TEST_F(HloReplicationAnalysisTest, + NonReplicatedConditionMarksBodyAtEveryIndex) { + const std::string module_str = R"hlo( +HloModule NonReplicatedConditionMarksBodyAtEveryIndex + +cond { + cond_param = (f32[8], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + limit = u32[] constant(5) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[8], u32[]) parameter(0) + x = f32[8] get-tuple-element(body_param), index=0 + i = u32[] get-tuple-element(body_param), index=1 + replica-id = u32[] replica-id() + next_i = u32[] add(i, replica-id) + ROOT tuple = (f32[8], u32[]) tuple(x, next_i) +} + +ENTRY entry { + p0 = f32[8] parameter(0), parameter_replication={true} + zero = u32[] constant(0) + init = (f32[8], u32[]) tuple(p0, zero) + ROOT while = (f32[8], u32[]) while(init), condition=cond, body=body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::Run( + module.get(), /*cross_partition_spmd=*/false)); + const HloInstruction* body_param = + FindInstruction(module.get(), "body_param"); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(body_param, {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(body_param, {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt(body_param, {1})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "x"), {})); + // The body root tuple is unique at its own index too, and so are the loop + // result and the condition parameter that merge it. + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "tuple"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "cond_param"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "cond_param"), {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while"), {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while"), {1})); + // The loop input is outside the body and keeps its replicated tuple root. + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "init"), {})); +} + +// The trip count depends on the partition, so the partitions run the infeed in +// the body a different number of times. Its data is unique although its +// sharding is replicated, and so is the loop result that it feeds. +TEST_F(HloReplicationAnalysisTest, InfeedInDivergentLoopIsNotReplicatedSPMD) { + const std::string module_str = R"hlo( +HloModule InfeedInDivergentLoopIsNotReplicatedSPMD + +cond { + cond_param = (f32[8], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + pid = u32[] partition-id() + cmp_val = u32[] add(i, pid) + limit = u32[] constant(4) + ROOT lt = pred[] compare(cmp_val, limit), direction=LT +} + +body { + body_param = (f32[8], u32[]) parameter(0) + tok = token[] after-all() + infeed = (f32[8], token[]) infeed(tok), + sharding={{replicated}, {maximal device=0}} + infeed_data = f32[8] get-tuple-element(infeed), index=0 + i = u32[] get-tuple-element(body_param), index=1 + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[8], u32[]) tuple(infeed_data, next_i) +} + +ENTRY entry { + p0 = f32[8] parameter(0), sharding={replicated} + zero = u32[] constant(0) + init = (f32[8], u32[]) tuple(p0, zero) + ROOT while = (f32[8], u32[]) while(init), condition=cond, body=body +} +)hlo"; + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(module_str, /*replica_count=*/1, + /*num_partitions=*/2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr analysis, + HloReplicationAnalysis::Run(module.get(), /*cross_partition_spmd=*/true)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "lt"), {})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "infeed"), {0})); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "infeed_data"), {})); + for (const char* name : {"tuple", "while"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {0})) + << name; + } + // The replicated parameter outside the loop stays replicated. + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "init"), {0})); +} + +// Two loops in a called computation share their body. The second loop makes +// the shared body unique at index 0 after the first loop took its result, and +// nothing that feeds the call changes. The next pass of the outer loop still +// has to visit the call and the first loop again for their results to become +// unique. +TEST_F(HloReplicationAnalysisTest, SharedLoopBodyChangesAfterItsFirstCaller) { + const std::string module_str = R"hlo( +HloModule SharedLoopBodyChangesAfterItsFirstCaller + +cond { + cond_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[], u32[]) parameter(0) + x = f32[] get-tuple-element(body_param), index=0 + i = u32[] get-tuple-element(body_param), index=1 + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[], u32[]) tuple(x, next_i) +} + +callee { + a = f32[] parameter(0) + b = f32[] parameter(1) + zero = u32[] constant(0) + init1 = (f32[], u32[]) tuple(a, zero) + while1 = (f32[], u32[]) while(init1), condition=cond, body=body + init2 = (f32[], u32[]) tuple(b, zero) + while2 = (f32[], u32[]) while(init2), condition=cond, body=body + r1 = f32[] get-tuple-element(while1), index=0 + r2 = f32[] get-tuple-element(while2), index=0 + ROOT out = (f32[], f32[]) tuple(r1, r2) +} + +outer_cond { + outer_cond_param = (f32[], f32[], u32[]) parameter(0) + j = u32[] get-tuple-element(outer_cond_param), index=2 + limit = u32[] constant(3) + ROOT lt = pred[] compare(j, limit), direction=LT +} + +outer_body { + outer_param = (f32[], f32[], u32[]) parameter(0) + outer_a = f32[] get-tuple-element(outer_param), index=0 + outer_b = f32[] get-tuple-element(outer_param), index=1 + outer_j = u32[] get-tuple-element(outer_param), index=2 + call = (f32[], f32[]) call(outer_a, outer_b), to_apply=callee + call_r1 = f32[] get-tuple-element(call), index=0 + one = u32[] constant(1) + next_j = u32[] add(outer_j, one) + ROOT outer_tuple = (f32[], f32[], u32[]) tuple(call_r1, outer_b, next_j) +} + +ENTRY entry { + p0 = f32[] parameter(0), parameter_replication={true} + p1 = f32[] parameter(1), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[], f32[], u32[]) tuple(p0, p1, zero) + ROOT outer = (f32[], f32[], u32[]) while(init), condition=outer_cond, + body=outer_body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::Run( + module.get(), /*cross_partition_spmd=*/false)); + for (const char* name : {"while1", "call", "outer"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {0})) + << name; + } + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "outer_a"), {})); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while1"), {1})); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "outer_j"), {})); +} + +// Two loops in different called computations share a body that returns its +// parameter, and nothing else. The second loop makes the body unique at index +// 0 after the first loop took its result. Only the body root changes, so the +// first loop and its call have to be visited again when the outer loop +// repeats. +TEST_F(HloReplicationAnalysisTest, IdentityLoopBodySharedAcrossCalls) { + const std::string module_str = R"hlo( +HloModule IdentityLoopBodySharedAcrossCalls + +cond1 { + cond1_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond1_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +cond2 { + cond2_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond2_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + ROOT body_param = (f32[], u32[]) parameter(0) +} + +callee1 { + a = f32[] parameter(0) + zero = u32[] constant(0) + init1 = (f32[], u32[]) tuple(a, zero) + while1 = (f32[], u32[]) while(init1), condition=cond1, body=body + ROOT r1 = f32[] get-tuple-element(while1), index=0 +} + +callee2 { + b = f32[] parameter(0) + zero = u32[] constant(0) + init2 = (f32[], u32[]) tuple(b, zero) + while2 = (f32[], u32[]) while(init2), condition=cond2, body=body + ROOT r2 = f32[] get-tuple-element(while2), index=0 +} + +outer_cond { + outer_cond_param = (f32[], f32[], u32[], f32[]) parameter(0) + j = u32[] get-tuple-element(outer_cond_param), index=2 + limit = u32[] constant(3) + ROOT lt = pred[] compare(j, limit), direction=LT +} + +outer_body { + outer_param = (f32[], f32[], u32[], f32[]) parameter(0) + outer_a = f32[] get-tuple-element(outer_param), index=0 + outer_b = f32[] get-tuple-element(outer_param), index=1 + outer_j = u32[] get-tuple-element(outer_param), index=2 + call1 = f32[] call(outer_a), to_apply=callee1 + one = u32[] constant(1) + next_j = u32[] add(outer_j, one) + call2 = f32[] call(outer_b), to_apply=callee2 + ROOT outer_tuple = (f32[], f32[], u32[], f32[]) tuple(call1, outer_b, next_j, + call2) +} + +ENTRY entry { + p0 = f32[] parameter(0), parameter_replication={true} + p1 = f32[] parameter(1), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[], f32[], u32[], f32[]) tuple(p0, p1, zero, p1) + ROOT outer = (f32[], f32[], u32[], f32[]) while(init), condition=outer_cond, + body=outer_body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::Run( + module.get(), /*cross_partition_spmd=*/false)); + for (const char* name : {"while1", "outer"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {0})) + << name; + } + for (const char* name : {"call1", "outer_a"}) { + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), name), {})) + << name; + } + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "while1"), {1})); +} + +// An all-reduce over groups of one device yields a partially replicated value +// whose device sets are all singletons. Merging it with itself gives unique, so +// in a loop body it turns unique when the body is visited again: with partial +// replication every visit evaluates every instruction. +TEST_F(HloReplicationAnalysisTest, SingletonGroupsInLoopBodyBecomeUnique) { + const std::string module_str = R"hlo( +HloModule SingletonGroupsInLoopBodyBecomeUnique + +sum { + a = f32[] parameter(0) + b = f32[] parameter(1) + ROOT add = f32[] add(a, b) +} + +cond { + cond_param = (f32[], u32[]) parameter(0) + i = u32[] get-tuple-element(cond_param), index=1 + limit = u32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +body { + body_param = (f32[], u32[]) parameter(0) + x = f32[] get-tuple-element(body_param), index=0 + i = u32[] get-tuple-element(body_param), index=1 + all-reduce = f32[] all-reduce(x), replica_groups={{0},{1}}, to_apply=sum + one = u32[] constant(1) + next_i = u32[] add(i, one) + ROOT tuple = (f32[], u32[]) tuple(all-reduce, next_i) +} + +ENTRY entry { + p0 = f32[] parameter(0), parameter_replication={false} + zero = u32[] constant(0) + init = (f32[], u32[]) tuple(p0, zero) + ROOT while = (f32[], u32[]) while(init), condition=cond, body=body +} +)hlo"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/2)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloReplicationAnalysis::RunWithPartialReplication( + module.get(), /*cross_partition_spmd=*/false)); + const std::vector singletons = CreateReplicaGroups({{0}, {1}}); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "all-reduce"), {}, singletons)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "tuple"), {0}, singletons)); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "tuple"), {1})); +} + } // namespace } // namespace xla From 7135e07461162270425f4bbfd225873636c37996 Mon Sep 17 00:00:00 2001 From: Sohaib Iftikhar Date: Tue, 29 Sep 2026 07:51:31 -0700 Subject: [PATCH 16/25] [XLA:GPU] Add BarrierMode attribute to BlockBarrierOp. We already had symmetric and kConsumerSymmetric where a single consumer waits for multiple producers on the same gpu. Now we also need producer symmetric for the case where a consumer waits for the same producer block from different GPUs. PiperOrigin-RevId: 990304779 --- .../gpu/codegen/triton/ir/tests/invalid.mlir | 24 ++++++++- .../gpu/codegen/triton/ir/triton_xla_attrs.td | 14 +++++ .../gpu/codegen/triton/ir/triton_xla_ops.cc | 19 ++++--- .../gpu/codegen/triton/ir/triton_xla_ops.td | 24 +++++++-- .../tests/triton_xla_block_barrier.mlir | 53 +++++++++++++++++-- .../triton_xla_lower_block_barrier_pass.cc | 32 +++++++---- 6 files changed, 139 insertions(+), 27 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir index b4ed1827ab6ff3..603a011b0af778 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir @@ -79,16 +79,36 @@ func.func @insert_wrong_shape(%arg0: tensor<8xbf16>, %arg1: !tt.ptr) { // ----- func.func @block_barrier_stride_without_slot(%ptr: !tt.ptr, %rank: i32, %sig: i32) { - // expected-error @+1 {{'triton_xla.block_barrier' op signal_stride requires signal_slot to be specified}} + // expected-error @+1 {{'triton_xla.block_barrier' op symmetric barrier_mode requires signal_slot to be omitted and signal_stride to be 0}} triton_xla.block_barrier %ptr, %rank, %sig : (!tt.ptr, i32, i32) -> () return } +func.func @block_barrier_symmetric_with_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op symmetric barrier_mode requires signal_slot to be omitted and signal_stride to be 0}} + triton_xla.block_barrier %ptr, %rank, %sig, : (!tt.ptr, i32, i32) -> () + return +} + // ----- + func.func @block_barrier_slot_without_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32, %slot: i32) { - // expected-error @+1 {{'triton_xla.block_barrier' op signal_slot requires a positive signal_stride}} + // expected-error @+1 {{'triton_xla.block_barrier' op symmetric barrier_mode requires signal_slot to be omitted and signal_stride to be 0}} triton_xla.block_barrier %ptr, %rank, %sig, %slot : (!tt.ptr, i32, i32, i32) -> () return } +func.func @block_barrier_asymmetric_without_slot(%ptr: !tt.ptr, %rank: i32, %sig: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op asymmetric barrier_mode requires signal_slot to be specified}} + triton_xla.block_barrier %ptr, %rank, %sig : (!tt.ptr, i32, i32) -> () + return +} + +// ----- + +func.func @block_barrier_asymmetric_without_stride(%ptr: !tt.ptr, %rank: i32, %sig: i32, %slot: i32) { + // expected-error @+1 {{'triton_xla.block_barrier' op asymmetric barrier_mode requires a positive signal_stride}} + triton_xla.block_barrier %ptr, %rank, %sig, %slot : (!tt.ptr, i32, i32, i32) -> () + return +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td index ab18493965e4d7..de54d77026f7da 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td @@ -86,5 +86,19 @@ def TTXLA_ComparatorAttr : EnumAttr let assemblyFormat = "$value"; } +def TTXLA_BarrierModeEnum : I32EnumAttr<"BarrierMode", + "Synchronization mode for triton_xla.block_barrier.", + [ + I32EnumAttrCase<"kSymmetric", 0, "symmetric">, + I32EnumAttrCase<"kConsumerSymmetric", 1, "consumer_symmetric">, + I32EnumAttrCase<"kProducerSymmetric", 2, "producer_symmetric">, + ]> { + let cppNamespace = "::mlir::triton::xla"; + let genSpecializedAttr = 0; +} +def TTXLA_BarrierModeAttr + : EnumAttr { + let assemblyFormat = "`<` $value `>`"; +} #endif // XLA_BACKENDS_GPU_CODEGEN_TRITON_IR_TRITON_XLA_ATTRS_TD_ diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc index 98ae71a0bb3c3f..608eecd7de12e2 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.cc @@ -324,14 +324,21 @@ LogicalResult PtrToMemrefOp::verify() { } LogicalResult BlockBarrierOp::verify() { - if (getSignalStride() < 0) { - return emitOpError("signal_stride must be non-negative"); + if (getBarrierMode() == BarrierMode::kSymmetric) { + if (getSignalSlot() || getSignalStride() != 0) { + return emitOpError( + "symmetric barrier_mode requires signal_slot to be omitted and " + "signal_stride to be 0"); + } + return success(); } - if (getSignalStride() > 0 && !getSignalSlot()) { - return emitOpError("signal_stride requires signal_slot to be specified"); + if (getSignalStride() <= 0) { + return emitOpError( + "asymmetric barrier_mode requires a positive signal_stride"); } - if (getSignalSlot() && getSignalStride() == 0) { - return emitOpError("signal_slot requires a positive signal_stride"); + if (!getSignalSlot()) { + return emitOpError( + "asymmetric barrier_mode requires signal_slot to be specified"); } return success(); } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td index 88aa3c2cf4a1e6..b8ebaee318e174 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td +++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td @@ -337,9 +337,9 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { of `i32` flags with shape `[num_blocks][world_size]`, flattened as `block_idx * world_size + slot_idx`. - This barrier supports two synchronization modes: + This barrier supports three synchronization modes via `$barrier_mode`: - ### 1. Symmetric Mode (`$signal_stride == 0`, `$signal_slot` omitted) + ### 1. `symmetric` (default; `$signal_stride == 0`, `$signal_slot` omitted) Used when block `block_idx` on every GPU produces and consumes data for the same `block_idx`: - **Signal (Write)**: Block `block_idx` on GPU `r` atomically writes @@ -349,7 +349,7 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { own buffer `$signal_buffers[r]` until all `world_size` slots at `block_idx * world_size + [0..world_size)` reach `>= $signal_value`. - ### 2. Asymmetric Mode (`$signal_stride > 0`, `$signal_slot` provided). + ### 2. `consumer_symmetric` (`$signal_stride > 0`, `$signal_slot` provided) Used when `world_size` blocks on the same GPU (spaced by `$signal_stride` in block ID) each produce `$signal_slot` (`0 <= $signal_slot < world_size`) of a tile that will be consumed by a single target block on every peer GPU: @@ -366,6 +366,19 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { slots at `block_idx * world_size + [0..world_size)` reach `>= $signal_value`. + ### 3. `producer_symmetric` (`$signal_stride > 0`, `$signal_slot` provided) + Used when block `block_idx` on every GPU produces data at its own + `block_idx`, and consumer block `block_idx` on GPU `r` waits for data + produced by `target_block` across all peer GPUs: + - Let `base_block = block_idx - $signal_slot * $signal_stride` and + `target_block = base_block + r * $signal_stride`. + - **Signal (Write)**: Block `block_idx` on GPU `r` atomically writes + `$signal_value` to index `block_idx * world_size + r` in every peer GPU's + signal buffer (`$signal_buffers[0..world_size)`). + - **Wait (Read)**: Block `block_idx` on GPU `r` spin-waits on its own + buffer `$signal_buffers[r]` until all `world_size` slots at + `target_block * world_size + [0..world_size)` reach `>= $signal_value`. + *NOTE* Block orchestration order is not guaranteed across GPUs, so there is a chance of a deadlock should the number of blocks be greater than the number of SMs on a GPU. That is to say that if BlockN is launched @@ -404,7 +417,10 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> { // same as this value. I32Attr:$world_size, // Optional block-ID stride between consumer blocks for asymmetric barriers. - DefaultValuedOptionalAttr:$signal_stride + DefaultValuedOptionalAttr:$signal_stride, + // Synchronization mode for the barrier. + DefaultValuedAttr:$barrier_mode ); // This is a barrier op and produces no results. diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir index 94c56e69205642..d587b6c3788d1f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir @@ -50,12 +50,12 @@ tt.func @block_barrier_kernel( tt.return } -// CHECK-LABEL: block_barrier_asymmetric_kernel +// CHECK-LABEL: block_barrier_consumer_symmetric_kernel // CHECK-SAME: %[[PTR:[^:]+]]: !tt.ptr, // CHECK-SAME: %[[SIGNAL_VALUE:[^:]+]]: i32, // CHECK-SAME: %[[RANK:[^:]+]]: i32, // CHECK-SAME: %[[SIGNAL_SLOT:[^:]+]]: i32 -tt.func @block_barrier_asymmetric_kernel( +tt.func @block_barrier_consumer_symmetric_kernel( %ptr : !tt.ptr, %signal_value : i32, %rank: i32, %signal_slot: i32) { // CHECK-NEXT: %[[STRIDE:.+]] = arith.constant 4 : i32 // CHECK-NEXT: %[[WORLD_SIZE:.+]] = arith.constant 2 : i32 @@ -88,7 +88,52 @@ tt.func @block_barrier_asymmetric_kernel( // CHECK-NEXT: } // CHECK-NEXT: ttg.barrier local // CHECK-NEXT: tt.return - triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot : + triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot + : (!tt.ptr, i32, i32, i32) -> () tt.return -} \ No newline at end of file +} + +// CHECK-LABEL: block_barrier_producer_symmetric_kernel +// CHECK-SAME: %[[PTR:[^:]+]]: !tt.ptr, +// CHECK-SAME: %[[SIGNAL_VALUE:[^:]+]]: i32, +// CHECK-SAME: %[[RANK:[^:]+]]: i32, +// CHECK-SAME: %[[SIGNAL_SLOT:[^:]+]]: i32 +tt.func @block_barrier_producer_symmetric_kernel( + %ptr : !tt.ptr, %signal_value : i32, %rank: i32, %signal_slot: i32) { + // CHECK-NEXT: %[[STRIDE:.+]] = arith.constant 4 : i32 + // CHECK-NEXT: %[[WORLD_SIZE:.+]] = arith.constant 2 : i32 + // CHECK-NEXT: %[[TID:.+]] = triton_xla.get_tid + // CHECK-NEXT: %[[PROGRAM_ID:.+]] = tt.get_program_id x + // CHECK-NEXT: %[[COND:.+]] = arith.cmpi ult, %[[TID]], %[[WORLD_SIZE]] + // CHECK-NEXT: scf.if %[[COND]] { + // CHECK-NEXT: %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]] + // CHECK-NEXT: %[[RANGE:.+]] = tt.make_range {end = 2 : i32, start = 0 : i32} + // CHECK-NEXT: %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]] + // CHECK-NEXT: %[[LOAD:.+]] = tt.load %[[ADD_PTR]] + // CHECK-NEXT: %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]] + // CHECK-NEXT: %[[MULI:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]] + // CHECK-NEXT: %[[S_OFFSET:.+]] = arith.muli %[[SIGNAL_SLOT]], %[[STRIDE]] + // CHECK-NEXT: %[[BASE_T:.+]] = arith.subi %[[PROGRAM_ID]], %[[S_OFFSET]] + // CHECK-NEXT: %[[RANK_OFFSET:.+]] = arith.muli %[[RANK]], %[[STRIDE]] + // CHECK-NEXT: %[[WAIT_BLOCK:.+]] = arith.addi %[[BASE_T]], %[[RANK_OFFSET]] + // CHECK-NEXT: %[[WAIT_BLOCK_OFFSET:.+]] = arith.muli %[[WAIT_BLOCK]], %[[WORLD_SIZE]] + // CHECK-NEXT: %[[WRITE_OFFSET:.+]] = arith.addi %[[MULI]], %[[RANK]] + // CHECK-NEXT: %[[SPLAT_WRITE_OFFSET:.+]] = tt.splat %[[WRITE_OFFSET]] + // CHECK-NEXT: %[[ADD_PTR_2:.+]] = tt.addptr %[[INT_TO_PTR]], %[[SPLAT_WRITE_OFFSET]] + // CHECK-NEXT: triton_xla.atomic_write sys, release, %[[ADD_PTR_2]], %[[SIGNAL_VALUE]] + // CHECK-NEXT: %[[ADD_PTR_3:.+]] = tt.addptr %[[PTR]], %[[RANK]] + // CHECK-NEXT: %[[LOAD_2:.+]] = tt.load %[[ADD_PTR_3]] + // CHECK-NEXT: %[[INT_TO_PTR_2:.+]] = tt.int_to_ptr %[[LOAD_2]] + // CHECK-NEXT: %[[ADD_PTR_4:.+]] = tt.addptr %[[INT_TO_PTR_2]], %[[WAIT_BLOCK_OFFSET]] + // CHECK-NEXT: %[[SPLAT_ADD_PTR_4:.+]] = tt.splat %[[ADD_PTR_4]] + // CHECK-NEXT: %[[ADD_PTR_5:.+]] = tt.addptr %[[SPLAT_ADD_PTR_4]], %[[RANGE]] + // CHECK-NEXT: triton_xla.atomic_spin_wait sys, acquire, %[[ADD_PTR_5]], greater_than_or_equal_to, %[[SIGNAL_VALUE]] + // CHECK-NEXT: } + // CHECK-NEXT: ttg.barrier local + // CHECK-NEXT: tt.return + triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot + : + (!tt.ptr, i32, i32, i32) -> () + tt.return +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc index 41f48e48010301..f53e2aa9a7f030 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc @@ -52,6 +52,7 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, block_barrier.getSignalValue(); const int32_t world_size = block_barrier.getWorldSize(); const int32_t signal_stride = block_barrier.getSignalStride(); + const BarrierMode barrier_mode = block_barrier.getBarrierMode(); // Triton magic constant. constexpr auto kGlobalAddressSpace = mlir::triton::PtrAddrSpace::Global; @@ -128,7 +129,8 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, auto block_offset = mlir::arith::MulIOp::create( builder, i32_type, block_id, world_size_op); mlir::Value write_offset; - if (signal_stride > 0) { + mlir::Value wait_block_offset = block_offset; + if (barrier_mode != BarrierMode::kSymmetric) { // Suppose we have 4 blocks (0, 1, 2, 3), 2 ranks (world_size = 2), // and signal_stride = 2. // @@ -144,12 +146,14 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, // - Block 1: base_block = 1 - (0 * 2) = 1 // - Block 3: base_block = 3 - (1 * 2) = 1 // - // 2. Find `target_block`, the block on peer GPUs that will read this - // rank's data (`base_block + rank * signal_stride`). - // - // 3. Write to `signal_slot` inside `target_block`'s slots - // (`target_block * world_size + signal_slot`), so blocks 1 and 3 - // land in adjacent slots 0 and 1 for `target_block`. + // 2. Find `target_block` (`base_block + rank * signal_stride`): + // - In `kConsumerSymmetric`, `target_block` is the block on peer + // GPUs that reads this rank's data, so we write to `signal_slot` + // inside `target_block`'s slots + // (`target_block * world_size + signal_slot`). + // - In `kProducerSymmetric`, `target_block` is the producer block + // whose slots (`target_block * world_size + [0..world_size)`) + // this consumer block waits on. mlir::Value signal_slot = block_barrier.getSignalSlot(); auto stride_op = mlir::arith::ConstantOp::create( builder, builder.getI32IntegerAttr(signal_stride)); @@ -163,8 +167,14 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, builder, i32_type, base_block, rank_offset); auto target_block_offset = mlir::arith::MulIOp::create( builder, i32_type, target_block, world_size_op); - write_offset = mlir::arith::AddIOp::create( - builder, i32_type, target_block_offset, signal_slot); + if (barrier_mode == BarrierMode::kConsumerSymmetric) { + write_offset = mlir::arith::AddIOp::create( + builder, i32_type, target_block_offset, signal_slot); + } else { + write_offset = mlir::arith::AddIOp::create(builder, i32_type, + block_offset, rank); + wait_block_offset = target_block_offset; + } } else { write_offset = mlir::arith::AddIOp::create(builder, i32_type, block_offset, rank); @@ -201,10 +211,10 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier, // -> !tt.ptr auto read_address = mlir::triton::IntToPtrOp::create( builder, ptr_to_i32_type, read_address_i64); - // Pointer to SignalBuffers[rank][block_id] + // Pointer to SignalBuffers[rank][wait_block_id] // -> !tt.ptr auto read_address_at_block_offset = mlir::triton::AddPtrOp::create( - builder, ptr_to_i32_type, read_address, block_offset); + builder, ptr_to_i32_type, read_address, wait_block_offset); // -> tensor> auto read_address_at_block_offset_tensor = mlir::triton::SplatOp::create(builder, tensor_of_ptr_to_i32_type, From e77517b46f0ca735b980d0297d2d0dbc1f28fb51 Mon Sep 17 00:00:00 2001 From: Sohaib Iftikhar Date: Tue, 29 Sep 2026 08:00:50 -0700 Subject: [PATCH 17/25] [XLA:GPU] Extract collective emitter helpers from AllReduceEmitter. Move some refactoring changes outside to reduce churn in the followup change. PiperOrigin-RevId: 990309442 --- .../gpu/codegen/triton/collective_emitter.cc | 150 +++++++++++------- 1 file changed, 91 insertions(+), 59 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc index 6f45f3bc155372..b6cdb9217cf98a 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc @@ -47,6 +47,7 @@ limitations under the License. #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Location.h" @@ -284,11 +285,14 @@ absl::Status ValidateBlockLevelFusionConfig( return absl::OkStatus(); } -absl::StatusOr> GetAllReduceUnmanagedKernelArguments( +// Unmanaged kernel arguments of collectives that use the AllReduce argument +// layout: rank, signal value, signal buffers and one remote buffer per +// parameter. +absl::StatusOr> GetRemoteBufferUnmanagedKernelArguments( const HloComputation* computation, - const HloAllReduceInstruction* all_reduce) { + const HloCollectiveInstruction* collective) { const int32_t num_devices = - all_reduce->device_list()->num_devices_per_group(); + collective->device_list()->num_devices_per_group(); std::vector unmanaged_arguments; unmanaged_arguments.reserve(computation->num_parameters() + kNumCollectiveMetadataArgs); @@ -421,6 +425,80 @@ mlir::Value EmitDeviceInvocationCount(mlir::ImplicitLocOpBuilder& b, b, counter, arith::ConstantOp::create(b, b.getI32IntegerAttr(1))); } +// Returns the element offset of the active half of a double buffer that holds +// `num_elements` elements of `element_type` per half. The lowest bit of +// `signal_value` selects the half. +mlir::Value EmitDoubleBufferOffset(mlir::ImplicitLocOpBuilder& b, + mlir::Value signal_value, + int64_t num_elements, + PrimitiveType element_type) { + const mlir::Type i64_type = b.getI64Type(); + mlir::Value signal_value_i64 = + arith::ExtSIOp::create(b, i64_type, signal_value); + mlir::Value one = + arith::ConstantOp::create(b, i64_type, b.getI64IntegerAttr(1)); + mlir::Value buffer_index = arith::AndIOp::create(b, signal_value_i64, one); + // The allocated double buffer size in bytes is aligned to + // kXlaAllocatedBufferAlignBytes. To get the offset to the second buffer, we + // divide the buffer size by 2 and then divide by the element size to get + // the number of elements in the buffer. Its important to do this to make + // sure that start of the buffers are always aligned to 16 bytes. + const int64_t element_size = ShapeUtil::ByteSizeOfPrimitiveType(element_type); + const int64_t buffer_size = xla::RoundUpTo( + num_elements * element_size, kXlaAllocatedBufferAlignBytes); + const int64_t elements_per_buffer = buffer_size / element_size; + mlir::Value elements_per_buffer_op = arith::ConstantOp::create( + b, i64_type, b.getI64IntegerAttr(elements_per_buffer)); + return arith::MulIOp::create(b, buffer_index, elements_per_buffer_op); +} + +// Emits instructions to get the pointer to the remote buffer of the given +// rank. +// We have a !tt.ptr pointing to the base of the remote +// buffers. We add the rank index to the base pointer and load to get to the +// base pointer of the remote buffer of the given rank. Then we add the buffer +// offset to get the pointer to the correct buffer inside (double buffering). +// Note that the returned pointer is of the storage type not the logical type. +mlir::Value EmitRemoteBufferPtr(mlir::ImplicitLocOpBuilder& b, + mlir::Value remote_buffers, mlir::Value rank, + mlir::Value buffer_offset, + mlir::Type elem_storage_type) { + const auto ptr_to_i64_type = + ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace); + const auto ptr_to_elem_type = + ttir::PointerType::get(elem_storage_type, kGlobalAddressSpace); + mlir::Value remote_buf_ptr_addr = + ttir::AddPtrOp::create(b, ptr_to_i64_type, remote_buffers, rank); + mlir::Value remote_buf_i64 = + ttir::LoadOp::create(b, // + remote_buf_ptr_addr, // + ttir::CacheModifier::NONE, // + ttir::EvictionPolicy::NORMAL, // + /*isVolatile=*/false); // + mlir::Value remote_buf_ptr_base = ttir::IntToPtrOp::create( + b, ptr_to_elem_type, remote_buf_i64, + llvm::ArrayRef{xtile::GetDivisibilityAttr(b)}); + return ttir::AddPtrOp::create(b, ptr_to_elem_type, remote_buf_ptr_base, + buffer_offset); +} + +// Emits a barrier between this block and the matching blocks of the other +// ranks. All threads in the block should have completed their writes before we +// proceed with a block barrier. +// Otherwise, remote ranks might start reading the data before it is ready. +void EmitBlockBarrier(mlir::ImplicitLocOpBuilder& b, mlir::Value signal_buffers, + mlir::Value rank, mlir::Value signal_value, + int64_t world_size, mlir::Value signal_slot = nullptr, + int64_t signal_stride = 0) { + mlir::triton::gpu::BarrierOp::create(b, mlir::triton::gpu::AddrSpace::Local); + mtx::BlockBarrierOp::create( + b, signal_buffers, rank, signal_value, signal_slot, + b.getI32IntegerAttr(static_cast(world_size)), + signal_stride > 0 + ? b.getI32IntegerAttr(static_cast(signal_stride)) + : mlir::IntegerAttr()); +} + class AllReduceEmitter { public: static mlir::LogicalResult Emit(AllReduceEmitterContext ctx, @@ -518,8 +596,6 @@ class AllReduceEmitter { remote_input_buffers_i64_ = ttir::BitcastOp::create( builder_, ptr_to_i64_type_, remote_input_buffers_); } - const mlir::Type i64_type = builder_.getI64Type(); - // Check if last bit of signal_value is 0 or 1. mlir::Value signal_value = signal_value_; // For two-shot signal value is always incremented by 2. // So we need to right shift it by 1 to get the buffer index. @@ -529,55 +605,18 @@ class AllReduceEmitter { arith::ConstantOp::create(builder_, signal_value.getType(), builder_.getI32IntegerAttr(1))); } - mlir::Value buffer_index = arith::AndIOp::create( - builder_, i64_type, - arith::ExtSIOp::create(builder_, i64_type, signal_value), // i32->i64 - arith::ConstantOp::create(builder_, i64_type, - builder_.getI64IntegerAttr(1))); - // The allocated double buffer size in bytes is aligned to - // kXlaAllocatedBufferAlignBytes. To get the offset to the second buffer, we - // divide the buffer size by 2 and then divide by the element size to get - // the number of elements in the buffer. Its important to do this to make - // sure that start of the buffers are always aligned to 16 bytes. - const int64_t buffer_size = xla::RoundUpTo( - ctx_.num_elements * - ShapeUtil::ByteSizeOfPrimitiveType(ctx_.element_type), - kXlaAllocatedBufferAlignBytes); - const int64_t elements_per_buffer = - buffer_size / ShapeUtil::ByteSizeOfPrimitiveType(ctx_.element_type); - buffer_offset_ = arith::MulIOp::create( - builder_, i64_type, buffer_index, - arith::ConstantOp::create( - builder_, i64_type, - builder_.getI64IntegerAttr(elements_per_buffer))); + buffer_offset_ = EmitDoubleBufferOffset( + builder_, signal_value, ctx_.num_elements, ctx_.element_type); initialized_ = true; return absl::OkStatus(); } - // Emits instructions to get the pointer to the remote buffer of the given - // rank. - // We have a !tt.ptr pointing to the base of the remote - // buffers. We add the rank index to the base pointer and load to get to the - // base pointer of the remote buffer of the given rank. Then we add the buffer - // offset to get the pointer to the correct buffer inside (double buffering). - // Note that the returned pointer is of the storage type not the logical type. + // Returns the pointer to the active remote buffer of the given rank. See + // EmitRemoteBufferPtr. mlir::Value GetRemoteBufferPtr(mlir::Value rank_idx) { CHECK(initialized_); - mlir::Value remote_buf_ptr_addr = ttir::AddPtrOp::create( - builder_, ptr_to_i64_type_, remote_input_buffers_i64_, rank_idx); - mlir::Value remote_buf_i64 = - ttir::LoadOp::create(builder_, // - remote_buf_ptr_addr, // - ttir::CacheModifier::NONE, // - ttir::EvictionPolicy::NORMAL, // - /*isVolatile=*/false); // - mlir::Value remote_buf_ptr_base = - ttir::IntToPtrOp::create(builder_, ptr_to_elem_type_, remote_buf_i64, - llvm::ArrayRef{ - xtile::GetDivisibilityAttr(builder_)}); - mlir::Value remote_buf_ptr = ttir::AddPtrOp::create( - builder_, ptr_to_elem_type_, remote_buf_ptr_base, buffer_offset_); - return remote_buf_ptr; + return EmitRemoteBufferPtr(builder_, remote_input_buffers_i64_, rank_idx, + buffer_offset_, elem_storage_type_); } // Loads a tile from the remote buffer of the given rank. @@ -674,15 +713,8 @@ class AllReduceEmitter { mlir::LogicalResult EmitSync(mlir::Value signal_value) { CHECK(initialized_); - // All threads in the block should have completed their writes before we - // proceed with a block barrier. - // Otherwise, remote ranks might start reading the data before it is ready. - mlir::triton::gpu::BarrierOp::create(builder_, - mlir::triton::gpu::AddrSpace::Local); - mtx::BlockBarrierOp::create(builder_, signal_buffers_, device_rank_, - signal_value, /*signal_slot=*/nullptr, - builder_.getI32IntegerAttr(ctx_.world_size), - /*signal_stride=*/nullptr); + EmitBlockBarrier(builder_, signal_buffers_, device_rank_, signal_value, + ctx_.world_size); return mlir::success(); } @@ -1154,8 +1186,8 @@ absl::StatusOr> GetCollectiveUnmanagedKernelArguments( const HloInstruction* root = computation->root_instruction(); switch (root->opcode()) { case HloOpcode::kAllReduce: - return GetAllReduceUnmanagedKernelArguments( - computation, Cast(root)); + return GetRemoteBufferUnmanagedKernelArguments( + computation, Cast(root)); case HloOpcode::kAllGather: { // AllGather only needs the 3 metadata args (rank, signal_value, // signal_buffers). No per-parameter scratch buffer args because the From 0091a150554e20ee17c6860f0a01399eead16769 Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Tue, 29 Sep 2026 08:21:47 -0700 Subject: [PATCH 18/25] [XLA:CPU] Use signed ops for lowering symbolic expr when necessary. PiperOrigin-RevId: 990319867 --- .../elementwise/strided_overhanging_tiles.hlo | 41 +++++ .../emitters/transforms/simplify_affine.cc | 151 ++++++++++++------ .../transforms/tests/simplify_affine.mlir | 29 +++- 3 files changed, 169 insertions(+), 52 deletions(-) create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo new file mode 100644 index 00000000000000..2cdfa33a4c0ab5 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_overhanging_tiles.hlo @@ -0,0 +1,41 @@ +// RUN: test_correctness %s + +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xla-simplify-arith \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +fusion { + p0 = f32[224] parameter(0) + ROOT slice = f32[112] slice(p0), slice={[0:224:2]} +} + +ENTRY main { + p0 = f32[224] parameter(0) + ROOT fusion = f32[112] fusion(p0), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["16"]}], + "num_ctas":"3" + } + } + } +} + +// NEW-VECTOR: #[[$UB_MAP:.*]] = #xla.indexing_map<"(pid, tid) -> (min(max(pid * -48 - tid * 16 + 112, 0), 16)), domain: pid in [0, 2], tid in [0, 2]"> +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 2 : index]} +// NEW-VECTOR: scf.for %[[TID:.*]] = +// NEW-VECTOR: %[[ALLOCA:.*]] = memref.alloca() : memref<16xf32> +// NEW-VECTOR: %[[UB:.*]] = xla.apply_indexing #[[$UB_MAP]](%[[PID]], %[[TID]]) +// NEW-VECTOR: scf.for %[[IV:.*]] = %{{.*}} to %[[UB]] +// NEW-VECTOR: memref.load +// NEW-VECTOR: memref.store %{{.*}}, %[[ALLOCA]][%[[IV]]] + +// NEW-VECTOR-LLVM-NOT: xla.apply_indexing +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM: llvm.intr.smax diff --git a/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc b/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc index 15f442ed6544e5..c3cd5b545fb632 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/simplify_affine.cc @@ -90,9 +90,45 @@ void CollectArgs(SymbolicExpr expr, SymbolicExprType type, ret.push_back(expr); } +// Returns true if the top-level operation of `expr` can be lowered to an +// unsigned arith op. Operands of `expr` are not checked. +bool IsUnsignedLoweringSupported(SymbolicExpr expr, + RangeEvaluator& range_evaluator) { + if (!expr.IsBinaryOp()) { + return true; + } + // Mod and div can be lowered if their LHS is >= 0 and their RHS is a + // constant positive number. + if (expr.GetType() == SymbolicExprType::kMod || + expr.GetType() == SymbolicExprType::kFloorDiv || + expr.GetType() == SymbolicExprType::kCeilDiv) { + if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS())) { + return false; + } + auto rhs_range = range_evaluator.ComputeExpressionRange(expr.GetRHS()); + if (!rhs_range.IsPoint() || rhs_range.lower <= 0) { + return false; + } + } + // Max and min can be lowered if both operands are >= 0. + if (expr.GetType() == SymbolicExprType::kMax || + expr.GetType() == SymbolicExprType::kMin) { + if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS()) || + !range_evaluator.IsAlwaysPositiveOrZero(expr.GetRHS())) { + return false; + } + } + return true; +} + struct ExpressionEvaluator { - ExpressionEvaluator(ImplicitLocOpBuilder& builder, ValueRange operands) - : builder(builder), operands(operands) { + // If `range_evaluator` is null, all mod, div, min and max operations are + // lowered to unsigned arith ops. The caller must ensure that this is correct, + // e.g. by checking `IsLoweringSupported`. Otherwise, signed arith ops are + // emitted for the operations whose unsigned lowering is not provably correct. + ExpressionEvaluator(ImplicitLocOpBuilder& builder, ValueRange operands, + RangeEvaluator* range_evaluator = nullptr) + : builder(builder), operands(operands), range_evaluator(range_evaluator) { for (int i = 0; i < operands.size(); ++i) { variable_distances.push_back(Distance(builder, operands[i])); } @@ -123,6 +159,7 @@ struct ExpressionEvaluator { ImplicitLocOpBuilder& builder; ValueRange operands; + RangeEvaluator* range_evaluator; SmallVector variable_distances; }; @@ -153,26 +190,56 @@ Value ExpressionEvaluator::EvaluateAddMul(SymbolicExpr expr) { } Value ExpressionEvaluator::EvaluateExpression(SymbolicExpr expr) { + bool use_signed = range_evaluator != nullptr && + !IsUnsignedLoweringSupported(expr, *range_evaluator); switch (expr.GetType()) { case SymbolicExprType::kAdd: return EvaluateAddMul(expr); case SymbolicExprType::kMul: return EvaluateAddMul(expr); - case SymbolicExprType::kMod: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kFloorDiv: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kCeilDiv: - return builder.create( - EvaluateExpression(expr.GetLHS()), EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kMax: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); - case SymbolicExprType::kMin: - return builder.create(EvaluateExpression(expr.GetLHS()), - EvaluateExpression(expr.GetRHS())); + case SymbolicExprType::kMod: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (!use_signed) { + return builder.create(lhs, rhs); + } + // SymbolicExpr uses floor semantics: lhs mod rhs has the sign of rhs. + Value quotient = builder.create(lhs, rhs); + return builder.create( + lhs, builder.create(quotient, rhs)); + } + case SymbolicExprType::kFloorDiv: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } + case SymbolicExprType::kCeilDiv: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } + case SymbolicExprType::kMax: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } + case SymbolicExprType::kMin: { + Value lhs = EvaluateExpression(expr.GetLHS()); + Value rhs = EvaluateExpression(expr.GetRHS()); + if (use_signed) { + return builder.create(lhs, rhs); + } + return builder.create(lhs, rhs); + } case SymbolicExprType::kConstant: return builder.create(expr.GetValue()); case SymbolicExprType::kVariable: @@ -182,32 +249,14 @@ Value ExpressionEvaluator::EvaluateExpression(SymbolicExpr expr) { } } +// Returns true if all operations in `expr` can be lowered to unsigned arith +// ops. bool IsLoweringSupported(SymbolicExpr expr, RangeEvaluator& range_evaluator) { if (!expr.IsBinaryOp()) { return true; } - // Mod and div can be lowered if their LHS is >= 0 and their RHS is a - // constant positive number. - if (expr.GetType() == SymbolicExprType::kMod || - expr.GetType() == SymbolicExprType::kFloorDiv || - expr.GetType() == SymbolicExprType::kCeilDiv) { - if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS())) { - return false; - } - auto rhs_range = range_evaluator.ComputeExpressionRange(expr.GetRHS()); - if (!rhs_range.IsPoint() || rhs_range.lower <= 0) { - return false; - } - } - // Max and min can be lowered if both operands are >= 0. - if (expr.GetType() == SymbolicExprType::kMax || - expr.GetType() == SymbolicExprType::kMin) { - if (!range_evaluator.IsAlwaysPositiveOrZero(expr.GetLHS()) || - !range_evaluator.IsAlwaysPositiveOrZero(expr.GetRHS())) { - return false; - } - } - return IsLoweringSupported(expr.GetLHS(), range_evaluator) && + return IsUnsignedLoweringSupported(expr, range_evaluator) && + IsLoweringSupported(expr.GetLHS(), range_evaluator) && IsLoweringSupported(expr.GetRHS(), range_evaluator); } @@ -274,23 +323,27 @@ struct RewriteApplyIndexingOp : OpRewritePattern { results.reserve(symbolic_map.GetNumResults()); for (unsigned i = 0; i < symbolic_map.GetNumResults(); ++i) { SymbolicExpr result_expr = symbolic_map.GetResult(i); - // If the expression cannot be lowered, we convert it to affine.apply, - // since it supports more expression types. if (IsLoweringSupported(result_expr, range_evaluator)) { results.push_back( ExpressionEvaluator(b, operands).EvaluateExpression(result_expr)); - } else { - // TODO: b/446856305 - Create a SymbolicApplyOp. For now, we convert the - // SymbolicMap back to AffineMap and fall back to AffineApplyOp. - AffineMap sub_map = SymbolicMapToAffineMap(symbolic_map.GetSubMap({i})); - if (!sub_map) { - return rewriter.notifyMatchFailure(op, - "cannot fallback to affine.apply"); - } + continue; + } + // If the expression cannot be lowered to unsigned ops, we convert it to + // affine.apply, since it supports more expression types. + // TODO: b/446856305 - Create a SymbolicApplyOp. For now, we convert the + // SymbolicMap back to AffineMap and fall back to AffineApplyOp. + if (AffineMap sub_map = + SymbolicMapToAffineMap(symbolic_map.GetSubMap({i}))) { results.push_back(b.create( sub_map, operands.take_front(symbolic_map.GetNumDims() + symbolic_map.GetNumSymbols()))); + continue; } + // affine.apply does not support min and max. Lower the expression to + // arith ops, using signed ops wherever unsigned ones are not provably + // correct. + results.push_back(ExpressionEvaluator(b, operands, &range_evaluator) + .EvaluateExpression(result_expr)); } rewriter.replaceOp(op, results); return mlir::success(); diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir index e6e3ffe69d62ab..8d781e6977d758 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir +++ b/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_affine.mlir @@ -260,15 +260,38 @@ func.func @lower_ceildiv_constant(%arg0: index) -> index { // ----- -func.func @unsafe_max(%arg0: index {xla.range = [-10 : index, 10 : index]}, %arg1: index) -> index { +func.func @signed_max(%arg0: index {xla.range = [-10 : index, 10 : index]}, %arg1: index) -> index { %0 = xla.apply_indexing #xla.indexing_map< "()[s0, s1] -> (max(s0, s1))," "domain: s0 in [-10, 10], s1 in [0, 5]">[%arg0, %arg1] return %0 : index } -// CHECK-LABEL: @unsafe_max -// CHECK: xla.apply_indexing +// CHECK-LABEL: @signed_max +// CHECK-SAME: (%[[ARG0:.*]]: index {{.*}}, %[[ARG1:.*]]: index) +// CHECK-NOT: xla.apply_indexing +// CHECK: %[[RET:.*]] = arith.maxsi %[[ARG0]], %[[ARG1]] +// CHECK-NEXT: return %[[RET]] + +// ----- + +// min/max can't be expressed with affine.apply, so possibly negative operands +// are lowered to signed arith ops. Operations with provably non-negative +// operands still use unsigned ops. +func.func @signed_clamp_of_floordiv(%arg0: index, %arg1: index) -> index { + %0 = xla.apply_indexing + #xla.indexing_map< + "(d0, d1) -> (min(max((-(d1 mod 3) - (d1 floordiv 3) * 24 - d0 * 13920 + 111317) floordiv 3, 0), 8))," + "domain: d0 in [0, 7], d1 in [0, 1739]">(%arg0, %arg1) + return %0 : index +} +// CHECK-LABEL: @signed_clamp_of_floordiv +// CHECK-NOT: xla.apply_indexing +// CHECK-NOT: affine.apply +// CHECK: %[[DIV:.*]] = arith.floordivsi +// CHECK: %[[MAX:.*]] = arith.maxsi %[[DIV]] +// CHECK: %[[MIN:.*]] = arith.minui %[[MAX]] +// CHECK-NEXT: return %[[MIN]] // ----- From e7481545620135170f629b9a6114be826ec7e1b9 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 29 Sep 2026 08:53:59 -0700 Subject: [PATCH 19/25] Reverts changelist 988304835 PiperOrigin-RevId: 990335835 --- .../service/spmd/shardy/sdy_round_trip/BUILD | 2 - .../sdy_round_trip/import_shardy_attrs.cc | 64 +--- .../shardy/sdy_round_trip/shard_map_export.cc | 127 +----- .../shardy/sdy_round_trip/shard_map_import.cc | 107 +----- .../stablehlo_round_trip/export_shardings.cc | 39 +- ...ound_trip_pipeline_manual_computation.mlir | 39 -- ..._trip_import_pipeline_hlo_sharding_v3.mlir | 24 -- ...trip_shard_map_export_hlo_sharding_v3.mlir | 285 -------------- ...y_round_trip_shard_map_import_failure.mlir | 35 -- ...trip_shard_map_import_hlo_sharding_v3.mlir | 362 ------------------ 10 files changed, 39 insertions(+), 1045 deletions(-) delete mode 100644 third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir delete mode 100644 third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD index 8b05f417940548..2ca1c537bfc8d0 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD @@ -74,8 +74,6 @@ cc_library( deps = [ "//xla:side_effect_util", "//xla/hlo/ir:hlo_sharding", - "//xla/hlo/ir:mesh_and_axis", - "//xla/hlo/ir:named_sharding", "//xla/hlo/parser:hlo_parser", "//xla/hlo/translate/hlo_to_mhlo:hlo_utils", "//xla/service/spmd/shardy:constants", diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc index 4b3f5b085e2f92..30228198eaa7e2 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc @@ -55,8 +55,6 @@ limitations under the License. #include "shardy/dialect/sdy/ir/utils.h" #include "stablehlo/dialect/StablehloOps.h" #include "xla/hlo/ir/hlo_sharding.h" -#include "xla/hlo/ir/mesh_and_axis.h" -#include "xla/hlo/ir/named_sharding.h" #include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h" #include "xla/service/spmd/shardy/constants.h" @@ -84,7 +82,6 @@ using ::mlir::func::FuncOp; using ::mlir::sdy::getSharding; using ::mlir::sdy::kShardingAttr; using ::mlir::sdy::kShardingRuleAttr; -using ::mlir::sdy::ManualAxesAttr; using ::mlir::sdy::MeshAttr; using ::mlir::sdy::OpShardingRuleAttr; using ::mlir::sdy::PropagationBarrierOp; @@ -209,37 +206,6 @@ bool handleFuncResultSharding( return anyChanged; } -ManualAxesAttr extractManualAxes(const xla::HloSharding& hloSharding, - mlir::MLIRContext* context) { - if (hloSharding.IsTuple()) { - // Manual axes are only attached to ranked tensors, so a leading token - // element can legitimately have none. Take the first element that does. - for (const xla::HloSharding& element : hloSharding.tuple_elements()) { - if (ManualAxesAttr manualAxes = extractManualAxes(element, context)) { - return manualAxes; - } - } - return nullptr; - } - if (hloSharding.UseNamedShardingLeaf()) { - const xla::NamedSharding& namedSharding = hloSharding.named_sharding(); - llvm::SmallVector manualAxes; - const xla::Mesh& mesh = namedSharding.mesh(); - for (const auto& axisRef : namedSharding.manual_axes()) { - // Silently dropping an axis here would yield a manual computation with - // fewer manual axes than it actually has, which miscompiles rather than - // fails, so treat an out-of-range index as a bug. - CHECK_LT(axisRef.mesh_axis_index(), mesh.axis_names().size()); - manualAxes.push_back(mlir::StringAttr::get( - context, mesh.axis_names()[axisRef.mesh_axis_index()])); - } - if (!manualAxes.empty()) { - return ManualAxesAttr::get(context, manualAxes); - } - } - return nullptr; -} - // The sharding information is in the `kXlaShardingAttr` attribute. void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) { for (auto [argNum, argType] : llvm::enumerate(funcOp.getArgumentTypes())) { @@ -306,38 +272,12 @@ void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) { StringRef targetName = customCallOp.getCallTargetName(); if (targetName == kShardingCustomCallTargetName || targetName == "X64Combine" || - targetName == kGlobalToLocalShapeCallTargetName || - targetName == kLocalToGlobalShapeCallTargetName || isPythonCallbackCustomCall(customCallOp)) { - HloSharding hloSharding = parseShardingFromString(shardingAttr); if (auto sdySharding = convertToSdySharding( - hloSharding, customCallOp->getResultTypes(), - customCallOp->getContext())) { + parseShardingFromString(shardingAttr), + customCallOp->getResultTypes(), customCallOp->getContext())) { customCallOp->setAttr(kShardingAttr, sdySharding); } - // Only `xla.sdy.GlobalToLocalShape` (when there are operands) and the - // manual computation CallOp (when there are no operands) may carry - // manual axes. - if (targetName == kGlobalToLocalShapeCallTargetName) { - if (auto manualAxesAttr = - extractManualAxes(hloSharding, customCallOp->getContext())) { - customCallOp->setAttr(kManualAxes, manualAxesAttr); - } - } - } - } else if (auto callOp = mlir::dyn_cast(op); - callOp && isManualComputation(callOp)) { - // The body call op carries the local out shardings and the manual axes. - // It is the only carrier of the manual axes when the manual computation - // has no operands, since there is no `xla.sdy.GlobalToLocalShape` then. - HloSharding hloSharding = parseShardingFromString(shardingAttr); - if (auto sdySharding = convertToSdySharding( - hloSharding, callOp->getResultTypes(), callOp->getContext())) { - callOp->setAttr(kShardingAttr, sdySharding); - } - if (auto manualAxesAttr = - extractManualAxes(hloSharding, callOp->getContext())) { - callOp->setAttr(kManualAxes, manualAxesAttr); } } diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc index cb87f6e39b80e0..6b15a224661c57 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc @@ -17,10 +17,7 @@ limitations under the License. #include -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" -#include "llvm/Support/CommandLine.h" #include "llvm/Support/ErrorHandling.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/AffineMap.h" @@ -43,7 +40,6 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "mlir/Support/TypeID.h" #include "mlir/Transforms/DialectConversion.h" -#include "shardy/dialect/sdy/ir/constants.h" #include "shardy/dialect/sdy/ir/dialect.h" #include "shardy/dialect/sdy/ir/utils.h" #include "stablehlo/dialect/StablehloOps.h" @@ -63,44 +59,14 @@ using ::mlir::func::FuncOp; namespace stablehlo = ::mlir::stablehlo; namespace sdy = ::mlir::sdy; -// Returns `shardings` with `manualAxes` stripped from each element, i.e. the -// local view of tensors that are sharded over those axes. -// -// A `NamedSharding` may not mention the same axis twice, so an op whose result -// is a local tensor carries the manual axes in its `xla.sdy.manual_axes` -// attribute and the remaining axes in its dimension shardings. -sdy::TensorShardingPerValueAttr localizeShardings( - sdy::TensorShardingPerValueAttr shardings, - mlir::ArrayRef manualAxes, mlir::MLIRContext* context) { - llvm::SmallVector localShardings; - localShardings.reserve(shardings.size()); - for (sdy::TensorShardingAttr sharding : shardings.getShardings()) { - localShardings.push_back(sdy::eraseManualAxes(sharding, manualAxes)); - } - return sdy::TensorShardingPerValueAttr::get(context, localShardings); -} - class SdyRoundTripShardMapExportPass : public mlir::PassWrapper> { public: MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SdyRoundTripShardMapExportPass) - SdyRoundTripShardMapExportPass() = default; - explicit SdyRoundTripShardMapExportPass(bool enableHloShardingV3) { - this->enableHloShardingV3 = enableHloShardingV3; - } - SdyRoundTripShardMapExportPass(const SdyRoundTripShardMapExportPass& other) - : mlir::PassWrapper>(other) { - this->enableHloShardingV3 = other.enableHloShardingV3.getValue(); - } - - Option enableHloShardingV3{ - *this, "enable-hlo-sharding-v3", - llvm::cl::desc("Whether to enable HloShardingV3 which is the mesh and " - "axis based sharding representation."), - llvm::cl::init(false)}; + SdyRoundTripShardMapExportPass(bool enableHloShardingV3) + : enableHloShardingV3(enableHloShardingV3) {} void runOnOperation() final { ModuleOp moduleOp = getOperation(); @@ -124,81 +90,28 @@ class SdyRoundTripShardMapExportPass rewriter.setInsertionPoint(manualComputation); mlir::ValueRange operands = manualComputation->getOperands(); if (!operands.empty()) { - llvm::SmallVector globalOperands(operands.begin(), - operands.end()); - sdy::TensorShardingPerValueAttr inShardings = - manualComputation.getInShardings(); - if (enableHloShardingV3) { - inShardings = sdy::inlineMesh(symbolTable, inShardings); - // `GlobalToLocalShape`'s result is the *local* tensor, so it can only - // carry the local in shardings. Materialize the reshard that - // `ManualComputationOp` performs implicitly on its operands as an - // explicit `@Sharding` custom call. This keeps the exported module - // sharding consistent - every op's sharding describes its own result - // - and lets the import side read the global in shardings straight - // off the operand instead of inferring them. - for (auto [globalOperand, inSharding] : - llvm::zip_equal(globalOperands, inShardings.getShardings())) { - auto shardingOp = stablehlo::CustomCallOp::create( - rewriter, loc, globalOperand.getType(), globalOperand); - shardingOp.setCallTargetName(kShardingCustomCallTargetName); - // Import reads the global in shardings back off these wrappers, so - // they have to survive the HLO round trip intact. Without a side - // effect a wrapper of a zero-element tensor is replaced by a - // constant (see `ZeroSizedHloElimination`, which likewise spares - // the side-effecting `xla.sdy.GlobalToLocalShape`), which would - // lose that operand's sharding. The side effect also stops CSE from - // merging wrappers that carry different shardings. - shardingOp.setHasSideEffect(true); - shardingOp->setAttr(mlir::sdy::kShardingAttr, - sdy::TensorShardingPerValueAttr::get( - moduleOp.getContext(), inSharding)); - globalOperand = shardingOp.getResult(0); - } - } auto globalToLocalShape = stablehlo::CustomCallOp::create( - rewriter, loc, manualCompBodyArgTypes, globalOperands); + rewriter, loc, manualCompBodyArgTypes, operands); globalToLocalShape.setCallTargetName(kGlobalToLocalShapeCallTargetName); // We mark `xla.sdy.GlobalToLocalShape` as side-effecting to avoid // CSE deduping it with another taking the same operands, as it would // ignore the frontend attributes that could be different. globalToLocalShape.setHasSideEffect(true); + sdy::TensorShardingPerValueAttr inShardings = + manualComputation.getInShardings(); if (enableHloShardingV3) { - globalToLocalShape->setAttr( - mlir::sdy::kShardingAttr, - localizeShardings(inShardings, manualComputation.getManualAxes(), - moduleOp.getContext())); - globalToLocalShape->setAttr(kManualAxes, - manualComputation.getManualAxesAttr()); - } else { - setFrontendAttribute(globalToLocalShape, kInShardings, inShardings); - setFrontendAttribute(globalToLocalShape, kManualAxes, - manualComputation.getManualAxesAttr()); + inShardings = sdy::inlineMesh(symbolTable, inShardings); } + setFrontendAttribute(globalToLocalShape, kInShardings, inShardings); + setFrontendAttribute(globalToLocalShape, kManualAxes, + manualComputation.getManualAxesAttr()); operands = globalToLocalShape->getResults(); } - sdy::TensorShardingPerValueAttr outShardings = - manualComputation.getOutShardings(); - if (enableHloShardingV3) { - outShardings = sdy::inlineMesh(symbolTable, outShardings); - } - auto callOp = CallOp::create(rewriter, loc, localResultTypes, funcName, operands); setFrontendAttribute(callOp, kXlaInlineableAttr, rewriter.getStringAttr("xla_late")); - if (enableHloShardingV3 && !outShardings.getShardings().empty()) { - // The call op produces the local tensors, so it carries the local view - // of the out shardings together with the manual axes. This is also the - // only carrier of the manual axes when the manual computation has no - // operands, i.e. when there is no `xla.sdy.GlobalToLocalShape`. - callOp->setAttr( - mlir::sdy::kShardingAttr, - localizeShardings(outShardings, manualComputation.getManualAxes(), - moduleOp.getContext())); - callOp->setAttr(kManualAxes, manualComputation.getManualAxesAttr()); - } mlir::ResultRange results = manualComputation->getResults(); if (!results.empty()) { @@ -209,13 +122,14 @@ class SdyRoundTripShardMapExportPass // We mark `xla.sdy.LocalToGlobalShape` as side-effecting to avoid // CSE removing it if it has no users. localToGlobalShape.setHasSideEffect(true); + sdy::TensorShardingPerValueAttr outShardings = + manualComputation.getOutShardings(); if (enableHloShardingV3) { - localToGlobalShape->setAttr(mlir::sdy::kShardingAttr, outShardings); - } else { - setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings); - setFrontendAttribute(localToGlobalShape, kManualAxes, - manualComputation.getManualAxesAttr()); + outShardings = sdy::inlineMesh(symbolTable, outShardings); } + setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings); + setFrontendAttribute(localToGlobalShape, kManualAxes, + manualComputation.getManualAxesAttr()); results = localToGlobalShape->getResults(); } sdy::inlineRegionAndConvertTerminatorOp( @@ -233,16 +147,15 @@ class SdyRoundTripShardMapExportPass "1. A separate function for the body of the `ManualComputationOp`." "2. A `CallOp` calling the function in #1, marked as not inlinable." "3. A pair of `CustomCallOp`s that change the shape of the " - " arguments/results." - " Under HloShardingV3 the shardings and manual axes are attached as " - "native attributes, and each operand is additionally wrapped in a " - "`@Sharding` custom call holding the global in sharding so that " - "every op's sharding describes its own result. Otherwise they are " - "saved as frontend attrs on the pair of `CustomCallOp`s."; + " arguments/results. They save the in/out shardings and manual " + " axes as frontend attrs."; } void getDependentDialects(mlir::DialectRegistry& registry) const final { registry.insert(); } + + private: + bool enableHloShardingV3 = false; }; } // namespace diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc index 29ea30dd68e693..2e2151e8d11f62 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc @@ -75,50 +75,6 @@ using ::mlir::stablehlo::CustomCallOp; namespace sdy = ::mlir::sdy; -// Recovers the global in shardings of a manual computation from the -// `@Sharding` custom calls that export wraps around each operand of -// `globalToLocalShape`, and removes those wrappers. -// -// Export materializes the reshard that `sdy.manual_computation` performs -// implicitly on entry as an explicit `@Sharding` op, so that every op's -// sharding describes its own result. Undoing it here makes the round trip an -// identity. -// -// Returns an empty vector, and leaves the IR untouched, if any operand is not -// wrapped - the shardings are only meaningful as a complete set. -SmallVector takeGlobalInShardings( - CustomCallOp globalToLocalShape, mlir::IRRewriter& rewriter) { - SmallVector globalInShardings; - SmallVector shardingOps; - for (mlir::Value globalOperand : globalToLocalShape.getOperands()) { - auto shardingOp = globalOperand.getDefiningOp(); - if (!shardingOp || - shardingOp.getCallTargetName() != kShardingCustomCallTargetName) { - return {}; - } - sdy::TensorShardingAttr sharding = - sdy::getSharding(shardingOp.getResult(0)); - if (!sharding) { - return {}; - } - globalInShardings.push_back(sharding); - shardingOps.push_back(shardingOp); - } - - for (auto [i, shardingOp] : llvm::enumerate(shardingOps)) { - globalToLocalShape->setOperand(i, shardingOp.getOperand(0)); - } - // The same wrapper can feed several operands, so only visit it once - after - // the first erase the handle is dangling. - llvm::SmallDenseSet seen; - for (CustomCallOp shardingOp : shardingOps) { - if (seen.insert(shardingOp).second && shardingOp.use_empty()) { - rewriter.eraseOp(shardingOp); - } - } - return globalInShardings; -} - mlir::LogicalResult rewriteManualComputation( CallOp callOp, mlir::IRRewriter& rewriter, const mlir::SymbolTable& symbolTable) { @@ -140,7 +96,7 @@ mlir::LogicalResult rewriteManualComputation( // `ManualComputationOp` differently depending on whether the original had // operands/results. CustomCallOp globalToLocalShape; - SmallVector operands(callOp.getOperands()); + mlir::ValueRange operands = callOp.getOperands(); if (!operands.empty()) { // An input to `sdy.manual_computation` can have a dimension of size 0 // (i.e. 0 num-elements), in which case, the corresponding result of @@ -158,8 +114,7 @@ mlir::LogicalResult rewriteManualComputation( globalToLocalShape = (*customCallResIt).getDefiningOp(); CHECK_EQ(globalToLocalShape.getCallTargetName(), kGlobalToLocalShapeCallTargetName); - operands.assign(globalToLocalShape->operand_begin(), - globalToLocalShape->operand_end()); + operands = globalToLocalShape->getOperands(); } mlir::TypeRange resultTypes = callOp->getResultTypes(); @@ -195,13 +150,10 @@ mlir::LogicalResult rewriteManualComputation( if (!customCallOp) { return; } - if (mlir::DictionaryAttr frontendAttrs = getFrontendAttrs(customCallOp)) { - if (hasKey(frontendAttrs, shardingAttrName)) { - shardings = parseStringAttr( - frontendAttrs, shardingAttrName); - } - if (manualAxes.empty() && hasKey(frontendAttrs, kManualAxes)) { + shardings = parseStringAttr( + frontendAttrs, shardingAttrName); + if (manualAxes.empty()) { manualAxes = parseStringAttr(frontendAttrs, kManualAxes); } @@ -210,45 +162,6 @@ mlir::LogicalResult rewriteManualComputation( setShardingAttrs(globalToLocalShape, inShardings, kInShardings); setShardingAttrs(localToGlobalShape, outShardings, kOutShardings); - - // Under HloShardingV3 the shardings and manual axes round trip as native - // attributes rather than frontend attributes, so read them directly. - if (outShardings.empty() && localToGlobalShape) { - // `LocalToGlobalShape` produces the global tensors, so its own sharding is - // the global out shardings. - if (auto sharding = - localToGlobalShape->getAttrOfType( - sdy::kShardingAttr)) { - outShardings = sharding; - } - } - - if (manualAxes.empty()) { - // A manual computation without operands has no `GlobalToLocalShape`, in - // which case the body call op is the only carrier of the manual axes. - for (Operation* op : - {globalToLocalShape.getOperation(), callOp.getOperation()}) { - auto axes = - op ? op->getAttrOfType(kManualAxes) : nullptr; - if (axes && !axes.empty()) { - manualAxes = axes; - break; - } - } - } - - if (inShardings.empty() && globalToLocalShape) { - SmallVector globalInShardings = - takeGlobalInShardings(globalToLocalShape, rewriter); - if (!globalInShardings.empty()) { - inShardings = - sdy::TensorShardingPerValueAttr::get(context, globalInShardings); - // `takeGlobalInShardings` rewired `globalToLocalShape` past the wrappers. - operands.assign(globalToLocalShape->operand_begin(), - globalToLocalShape->operand_end()); - } - } - auto manualComputationOp = rewriter.replaceOpWithNewOp( callOp, resultTypes, operands, inShardings, outShardings, manualAxes); @@ -363,12 +276,10 @@ class SdyRoundTripShardMapImportPass } StringRef getDescription() const override { - return "converts a CallOp calling a @xla.sdy.manual_computation_body func, " - "wrapped with a pair of `CustomCallOps` that change the shape of " - "the arguments/results, to a ManualComputationOp. The in/out " - "shardings and manual axes are read from frontend attrs, or, under " - "HloShardingV3, from native attributes on the boundary ops and on " - "the `@Sharding` custom calls wrapping the operands"; + return "converts a CallOp calling a @xla.sdy.manual_computation_body func " + "with in/out shardings and manual axes as frontend attrs, wrapped " + "with a pair of `CustomCallOps` that change the shape of the " + "arguments/results, to a ManualComputationOp"; } void getDependentDialects(mlir::DialectRegistry& registry) const final { registry.insert(); diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc index e1f91c750f6882..0c76e727f68379 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc @@ -136,11 +136,6 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable, if (ManualAxesAttr manualAxesAttr = funcOp.getArgAttrOfType(argNum, kManualAxes)) { manualAxes = manualAxesAttr.getValue(); - if (mlir::isa( - funcOp.getArgument(argNum).getType()) && - getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { - manualAxes = {}; - } attrs.erase(kManualAxes); } attrs.set(kXlaShardingAttr, @@ -166,10 +161,6 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable, if (auto manualAxesAttr = mlir::dyn_cast_or_null(attrs.get(kManualAxes))) { manualAxes = manualAxesAttr.getValue(); - if (mlir::isa(funcOp.getResultTypes()[resNum]) && - getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { - manualAxes = {}; - } attrs.erase(kManualAxes); } attrs.set(kXlaShardingAttr, @@ -328,38 +319,24 @@ HloSharding getHloShardingForOp( std::function getMeshAttr, ArrayRef manualAxes, bool enableHloShardingV3, bool simplifyReplicatedShardings) { - auto getManualAxesForType = [&](mlir::Type type, - TensorShardingAttr sdySharding) { - if (mlir::isa(type) && - getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { - return ArrayRef(); - } - return manualAxes; - }; - bool isNoResultMaximal = op->getNumResults() == 0 && shardings.size() == 1 && (getMeshAttr(shardings.front()).isMaximal() || shardings.front().isFullyReplicated()); CHECK(shardings.size() == op->getNumResults() || isNoResultMaximal); if (op->getNumResults() == 1 || isNoResultMaximal) { - ArrayRef resultManualAxes = - op->getNumResults() == 1 - ? getManualAxesForType(op->getResultTypes().front(), - shardings.front()) - : manualAxes; - return convertToHloSharding(shardings.front(), getMeshAttr, - resultManualAxes, enableHloShardingV3, + return convertToHloSharding(shardings.front(), getMeshAttr, manualAxes, + enableHloShardingV3, simplifyReplicatedShardings); } std::vector newShardings; newShardings.reserve(shardings.size()); - for (auto [type, sdySharding] : - llvm::zip_equal(op->getResultTypes(), shardings)) { - newShardings.push_back(convertToHloSharding( - sdySharding, getMeshAttr, getManualAxesForType(type, sdySharding), - enableHloShardingV3, simplifyReplicatedShardings)); - } + llvm::transform(shardings, std::back_inserter(newShardings), + [&](TensorShardingAttr sdySharding) { + return convertToHloSharding(sdySharding, getMeshAttr, + manualAxes, enableHloShardingV3, + simplifyReplicatedShardings); + }); std::vector shapes; shapes.reserve(op->getNumResults()); diff --git a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir index 597d749388fdf6..9a05f216ae8a8e 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir @@ -154,42 +154,3 @@ func.func @main( } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> } - -// ----- - -// ***************** Zero-element operand test ***************** - -// A zero-element operand is replaced by a constant during the HLO round trip, -// which drops any non side-effecting op producing it. Make sure the remaining -// operands keep their in shardings rather than the manual computation coming -// back with none at all. - -// Make sure this temp attr doesn't exist anymore. -// CHECK-NOT: sharding_hlo_string - -// CHECK-V2: sdy.mesh @mesh_1 = <["a"=4]> -// CHECK-V3: sdy.mesh @mesh = <["a"=4]> -sdy.mesh @mesh_1 = <["a"=4]> - -// CHECK-LABEL: func.func @main -func.func @main(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> { - // CHECK-NEXT: %[[MANUAL_COMP:.*]] = sdy.manual_computation(%arg0) - // CHECK-V2-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) { - // CHECK-V3-SAME{LITERAL}: in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<1x8xi32> - // CHECK-NEXT: } : (tensor<4x8xi32>) -> tensor<4x8xi32> - // CHECK-NEXT: return %[[MANUAL_COMP]] : tensor<4x8xi32> - %c = stablehlo.constant dense<> : tensor<0xi32> - %0 = sdy.manual_computation(%c, %arg0) - in_shardings=[<@mesh_1, [{}], replicated={"a"}>, <@mesh_1, [{"a"}, {}]>] - out_shardings=[<@mesh_1, [{"a"}, {}]>] - manual_axes={"a"} (%arg1: tensor<0xi32>, %arg2: tensor<1x8xi32>) { - %init = stablehlo.constant dense<0> : tensor - %red = stablehlo.reduce(%arg1 init: %init) applies stablehlo.add across dimensions = [0] : (tensor<0xi32>, tensor) -> tensor - %bcast = stablehlo.broadcast_in_dim %red, dims = [] : (tensor) -> tensor<1x8xi32> - %add = stablehlo.add %arg2, %bcast : tensor<1x8xi32> - sdy.return %add : tensor<1x8xi32> - } : (tensor<0xi32>, tensor<4x8xi32>) -> tensor<4x8xi32> - return %0 : tensor<4x8xi32> -} - diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir index 1b0f94f6667750..58b7012ab5b707 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir @@ -204,30 +204,6 @@ module @module_1 { return %3 : tensor<1xi32> } - // Tests importing shard map custom calls annotated with HloShardingV3 mhlo.sharding. - // CHECK-LABEL: func @manual_computation_hlo_sharding_v3 - // CHECK-SAME: (%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> { - // CHECK-NEXT: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_2, [{"a"}, {}]>] out_shardings=[<@mesh_2, [{"a"}, {}]>] manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x16xf32>) { - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg1, %arg1 : tensor<2x16xf32> - // CHECK-NEXT: sdy.return %[[ADD]] : tensor<2x16xf32> - // CHECK-NEXT: } : (tensor<8x16xf32>) -> tensor<8x16xf32> - // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x16xf32> - // CHECK-NEXT: } - func.func @manual_computation_hlo_sharding_v3(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> { - %s0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<8x16xf32>) -> tensor<8x16xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{}, {}], manual={'a'}}"} : (tensor<8x16xf32>) -> tensor<2x16xf32> - %1 = call @xla.sdy.manual_computation_body_hlo_v3(%0) : (tensor<2x16xf32>) -> tensor<2x16xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<2x16xf32>) -> tensor<8x16xf32> - return %2 : tensor<8x16xf32> - } - - func.func private @xla.sdy.manual_computation_body_hlo_v3(%arg0: tensor<2x16xf32>) -> tensor<2x16xf32> { - %0 = stablehlo.add %arg0, %arg0 : tensor<2x16xf32> - return %0 : tensor<2x16xf32> - } - // CHECK-LABEL: func @frontend_attr_not_sharding // CHECK-SAME: %arg0: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", "b"}, {?}]>}, // CHECK-SAME: %arg1: tensor<16x8xf32> {mhlo.frontend_attributes = {baz = 1 : i32, foo = "bar"}}, diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir deleted file mode 100644 index fc83e872e68335..00000000000000 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir +++ /dev/null @@ -1,285 +0,0 @@ -// Copyright 2026 The OpenXLA 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. -// 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. -// ============================================================================== -// RUN: sdy_opt %s -xla-sdy-round-trip-shard-map-export='enable-hlo-sharding-v3=true' 2>&1 | FileCheck %s - -// Tests exporting `sdy.manual_computation` under HloShardingV3 (`enable-hlo-sharding-v3=true`). -// Verifies that: -// 1. Every op carries a native `sdy.sharding` that describes its own result: -// - a `@Sharding` custom call per operand holds the global in-shardings, -// making the reshard `sdy.manual_computation` performs implicitly on entry -// explicit. It is side-effecting so that the HLO round trip cannot drop it -// (see `ZeroSizedHloElimination` for a zero-element operand) or CSE two -// wrappers carrying different shardings; -// - `GlobalToLocalShape` and the body call op hold the local shardings plus -// typed `xla.sdy.manual_axes`; -// - `LocalToGlobalShape` holds the global out-shardings. -// 2. Legacy string `mhlo.frontend_attributes` (xla.sdy.in_shardings, etc.) are NOT emitted. - -sdy.mesh @mesh_0 = <["a"=4, "b"=2]> -sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]> - -// Ensure string frontend attributes are completely absent in V3 export: -// CHECK-NOT: xla.sdy.in_shardings -// CHECK-NOT: xla.sdy.out_shardings - -// Tests basic export with multiple inputs, outputs, and all-reduce. -// CHECK-LABEL: func @single_manual_comp -func.func @single_manual_comp(%arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", ?}, {"b", ?}]>}, %arg1: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"b", ?}, {?}]>}) -> (tensor<8x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) { - // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME: has_side_effect = true - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {"b"}]>]> - // CHECK-SAME: (tensor<8x16xf32>) -> tensor<8x16xf32> - // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1) - // CHECK-SAME: has_side_effect = true - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}, {}], replicated={"a"}>]> - // CHECK-SAME: (tensor<16x32xf32>) -> tensor<16x32xf32> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>, , [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>) - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}], replicated={"b"}>]> - // CHECK-SAME: (tensor<2x32xf32>) -> tensor<8x32xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<8x32xf32> - %0 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] manual_axes={"a", "b"} (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) { - %1 = stablehlo.add %arg2, %arg2 : tensor<2x8xf32> - %2 = stablehlo.dot %1, %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi64>, use_global_device_ids}> ({ - ^bb0(%arg4: tensor, %arg5: tensor): - %4 = stablehlo.add %arg4, %arg5 : tensor - stablehlo.return %4 : tensor - }) : (tensor<2x32xf32>) -> tensor<2x32xf32> - sdy.return %3 : tensor<2x32xf32> - } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> - return %0 : tensor<8x32xf32> -} - -// Tests chained manual computations where the output of one feeds the input of -// another with a *different* sharding. The `@Sharding` custom call on the -// second boundary is what preserves the reshard between them; without it the -// `in_shardings` of the second manual computation would be lost. -// CHECK-LABEL: func @manual_comp_using_another -func.func @manual_comp_using_another(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) - -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{}, {"b"}]>}) { - // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_0:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[SHMAP_0:.*]] = call @xla.sdy.manual_computation_body_0(%[[GLOBAL_TO_LOCAL_0]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_0:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_0]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%[[LOCAL_TO_GLOBAL_0]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_1:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x8xf32>) -> tensor<8x4xf32> - // CHECK-NEXT: %[[SHMAP_1:.*]] = call @xla.sdy.manual_computation_body_1(%[[GLOBAL_TO_LOCAL_1]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<8x4xf32>) -> tensor<8x4xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_1:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> - // CHECK-SAME: (tensor<8x4xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL_1]] : tensor<8x8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"a"}, {}]>] out_shardings=[<@mesh_0, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { - sdy.return %arg1 : tensor<2x8xf32> - } : (tensor<8x8xf32>) -> tensor<8x8xf32> - - %1 = sdy.manual_computation(%0) in_shardings=[<@mesh_0, [{}, {"b"}]>] out_shardings=[<@mesh_0, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x4xf32>) { - sdy.return %arg1 : tensor<8x4xf32> - } : (tensor<8x8xf32>) -> tensor<8x8xf32> - return %1 : tensor<8x8xf32> -} - -// Tests nested manual computations: an inner shard_map inside an outer shard_map body. -// CHECK-LABEL: func @nested_shmaps -func.func @nested_shmaps(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) { - // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<4x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_3(%[[GLOBAL_TO_LOCAL]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { - %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) { - %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - sdy.return %2 : tensor<2x4xf32> - } : (tensor<2x8xf32>) -> tensor<2x8xf32> - sdy.return %1 : tensor<2x8xf32> - } : (tensor<4x8xf32>) -> tensor<4x8xf32> - return %0 : tensor<4x8xf32> -} - -// Tests nested manual computation where the outer body performs additional math on the inner result. -// CHECK-LABEL: func @nested_shmaps_extra_op -func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) { - // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<4x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_5(%[[GLOBAL_TO_LOCAL]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> - // CHECK-SAME: (tensor<2x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32> - %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { - %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) { - %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - sdy.return %2 : tensor<2x4xf32> - } : (tensor<2x8xf32>) -> tensor<2x8xf32> - %3 = stablehlo.add %1, %1 : tensor<2x8xf32> - sdy.return %3 : tensor<2x8xf32> - } : (tensor<4x8xf32>) -> tensor<4x8xf32> - return %0 : tensor<4x8xf32> -} - -// Edge case: Manual computation producing an output but taking no inputs. There -// is no `GlobalToLocalShape`, so the body call op is the only carrier of the -// manual axes. -// CHECK-LABEL: func @manual_computation_no_inputs -func.func @manual_computation_no_inputs() -> tensor<4xi64> { - // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_6() - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: () -> tensor<2xi64> - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> - // CHECK-SAME: (tensor<2xi64>) -> tensor<4xi64> - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4xi64> - %0 = sdy.manual_computation() in_shardings=[] out_shardings=[<@mesh_0, [{"b"}]>] manual_axes={"b"} () { - %1 = stablehlo.constant dense<[2, 3]> : tensor<2xi64> - sdy.return %1 : tensor<2xi64> - } : () -> tensor<4xi64> - func.return %0 : tensor<4xi64> -} - -// Edge case: Manual computation consuming an input but producing no outputs. -// The body call op has no results, so `GlobalToLocalShape` is the only carrier -// of the manual axes. -// CHECK-LABEL: func @manual_computation_no_outputs -func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) { - // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (tensor<4xi64>) -> tensor<2xi64> - // CHECK-NEXT: call @xla.sdy.manual_computation_body_7(%[[GLOBAL_TO_LOCAL]]) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"}} - // CHECK-SAME: (tensor<2xi64>) -> () - // CHECK-NEXT: return - sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"b"}]>] out_shardings=[] manual_axes={"b"} (%arg1: tensor<2xi64>) { - stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> () - sdy.return - } : (tensor<4xi64>) -> () - func.return -} - -// Edge case: Degenerate manual computation with no inputs and no outputs. -// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs -func.func @manual_computation_no_inputs_no_outputs() { - // CHECK-NEXT: call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "xla_late"}} : () -> () - sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () { - sdy.return - } : () -> () - func.return -} - -// Type preservation: Token arguments should not have manual axes attached in V3. -// CHECK-LABEL: func @manual_computation_tokens -func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) { - // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>]> - // CHECK-SAME: (!stablehlo.token) -> !stablehlo.token - // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> - // CHECK-SAME: (tensor<2xi64>) -> tensor<2xi64> - // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]]) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>, , [{}]>]> - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>) - // CHECK-NEXT: %[[SHMAP:.*]]:2 = call @xla.sdy.manual_computation_body_9(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1) - // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} - // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy - // CHECK-SAME: (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) - // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]]:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]#0, %[[SHMAP]]#1) - // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>, , [{"b"}]>]> - // CHECK-SAME: (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>) - // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]]#0, %[[LOCAL_TO_GLOBAL]]#1 : !stablehlo.token, tensor<2xi64> - %0:2 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] manual_axes={"b"} (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) { - sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64> - } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) - func.return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> -} - - -// CHECK-LABEL: func @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_2(%[[GLOBAL_TO_LOCAL]]) -// CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32> - -// CHECK-LABEL: func @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_4(%[[GLOBAL_TO_LOCAL]]) -// CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} -// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x4xf32> -// CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) -// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> -// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x8xf32> -// CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[LOCAL_TO_GLOBAL]], %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32> -// CHECK-NEXT: return %[[ADD]] : tensor<2x8xf32> diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir index 15b7eec24bdca0..47d1ba107a8bf5 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir @@ -45,38 +45,3 @@ func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8 return %arg0 : tensor<0x8xf32> } -// ----- - -sdy.mesh @mesh = <["a"=2, "b"=2]> - -// Tests import error when manual computation body call operand is not produced by GlobalToLocalShape. -func.func @manual_computation_missing_global_to_local_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x16xf32>) { - %c = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - // expected-error @+2 {{'func.call' op expected at least one operand of @xla.sdy.manual_computation_body to be produced by a xla.sdy.GlobalToLocalShape CustomCallOp}} - // expected-error @+1 {{failed to rewrite func.call to manual computation}} - %0 = call @xla.sdy.manual_computation_body(%c) : (tensor<0x8xf32>) -> tensor<0x8xf32> - %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x8xf32>) -> (tensor<0x16xf32>) - return %1 : tensor<0x16xf32> -} - -func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> { - return %arg0 : tensor<0x8xf32> -} - -// ----- - -sdy.mesh @mesh = <["a"=2, "b"=2]> - -// Tests import error when manual computation body call result is not consumed by LocalToGlobalShape. -func.func @manual_computation_missing_local_to_global_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x8xf32>) { - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>) -> tensor<0x8xf32> - // expected-error @+2 {{'func.call' op expected the first use of @xla.sdy.manual_computation_body to be a xla.sdy.LocalToGlobalShape CustomCallOp}} - // expected-error @+1 {{failed to rewrite func.call to manual computation}} - %1 = call @xla.sdy.manual_computation_body(%0) : (tensor<0x8xf32>) -> tensor<0x8xf32> - return %1 : tensor<0x8xf32> -} - -func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> { - return %arg0 : tensor<0x8xf32> -} - diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir deleted file mode 100644 index 6e7060cb48ff30..00000000000000 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir +++ /dev/null @@ -1,362 +0,0 @@ -// Copyright 2026 The OpenXLA 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. -// 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. -// ============================================================================== -// RUN: sdy_opt %s --split-input-file -xla-sdy-test-flatten-call-graph -xla-sdy-round-trip-shard-map-import 2>&1 | FileCheck %s - -// Tests importing shard map custom calls back into `sdy.manual_computation` under -// HloShardingV3 (where frontend_attributes are omitted). -// Verifies that: -// 1. Shardings and manual axes are read directly off the native attributes of the -// boundary ops, without any inference: -// - global `in_shardings` come from the `@Sharding` custom call wrapping each -// operand of `GlobalToLocalShape`, which is then unwired and erased; -// - global `out_shardings` come from `LocalToGlobalShape`; -// - `manual_axes` come from `GlobalToLocalShape`, or from the body call op -// when the manual computation has no operands. -// 2. Outlined manual computation call bodies are inlined back into `sdy.manual_computation`. - -sdy.mesh @mesh_0 = <["a"=4, "b"=2]> -sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]> - -// Tests single manual computation with multiple inputs/outputs, all-reduce, and -// in/out shardings read off the `@Sharding` wrappers and `LocalToGlobalShape`. -// CHECK-LABEL: func @single_manual_comp -func.func @single_manual_comp(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<8x32xf32>) { - // CHECK-NOT: call @xla.sdy.manual_computation_body - // CHECK: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0, %arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] - - // CHECK-SAME{LITERAL}: manual_axes={"a", "b"} - // CHECK-SAME: (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) { - // CHECK-NEXT: %[[ADD_0:.*]] = stablehlo.add %arg2, %arg2 : tensor<2x8xf32> - // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %[[ADD_0]], %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - // CHECK-NEXT: %[[REDUCE:.*]] = "stablehlo.all_reduce"(%[[DOT]]) - // CHECK-NEXT: ^bb0(%arg4: tensor, %arg5: tensor): - // CHECK-NEXT: %[[ADD_1:.*]] = stablehlo.add %arg4, %arg5 : tensor - // CHECK-NEXT: stablehlo.return %[[ADD_1]] : tensor - // CHECK-NEXT: }) : (tensor<2x32xf32>) -> tensor<2x32xf32> - // CHECK-NEXT: sdy.return %[[REDUCE]] : tensor<2x32xf32> - // CHECK-NEXT: } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> - // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x32xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {"b"}]>]>} : (tensor<8x16xf32>) -> tensor<8x16xf32> - %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}], replicated={"a"}>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32> - %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>) - %1 = call @xla.sdy.manual_computation_body(%0#0, %0#1) : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}], replicated={"b"}>]>} : (tensor<2x32xf32>) -> tensor<8x32xf32> - return %2 : tensor<8x32xf32> -} - -// CHECK-LABEL: func @single_manual_comp_name_is_not_prefix_nor_suffix -func.func @single_manual_comp_name_is_not_prefix_nor_suffix(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>) { - // CHECK-NOT: call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234 - // CHECK: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<2x8xf32> - %1 = call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32> - return %2 : tensor<8x8xf32> -} - -// Tests chained manual computations: outputs of first shard map feed second shard map. -// CHECK-LABEL: func @manual_comp_using_another -func.func @manual_comp_using_another(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { - // CHECK-NOT: call @xla.sdy.manual_computation_body_0 - // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> - // CHECK-NOT: call @xla.sdy.manual_computation_body_1 - // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%[[MAN_COMP_0]]) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg1: tensor<8x4xf32>) { - // CHECK-NEXT: sdy.return %arg1 : tensor<8x4xf32> - // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> - // CHECK-NEXT: return %[[MAN_COMP_1]] : tensor<8x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<2x8xf32> - %1 = call @xla.sdy.manual_computation_body_0(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32> - %s1 = stablehlo.custom_call @Sharding(%2) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> - %3 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<8x4xf32> - %4 = call @xla.sdy.manual_computation_body_1(%3) : (tensor<8x4xf32>) -> tensor<8x4xf32> - %5 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%4) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x4xf32>) -> tensor<8x8xf32> - return %5 : tensor<8x8xf32> -} - -// Tests nested manual computations: reconstructing inner and outer sdy.manual_computation regions. -// CHECK-LABEL: func @nested_shmaps -func.func @nested_shmaps(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { - // CHECK-NOT: call @xla.sdy.manual_computation_body_2 - // CHECK-NOT: call @xla.sdy.manual_computation_body_3 - // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: tensor<2x4xf32>) { - // CHECK-NEXT: %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - // CHECK-NEXT: sdy.return %[[MUL]] : tensor<2x4xf32> - // CHECK-NEXT: } : (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: sdy.return %[[MAN_COMP_1]] : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<4x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[MAN_COMP_0]] : tensor<4x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4x8xf32>) -> tensor<2x8xf32> - %1 = call @xla.sdy.manual_computation_body_3(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32> - return %2 : tensor<4x8xf32> -} - -// Tests nested manual computation with operations in the outer region around the inner region. -// CHECK-LABEL: func @nested_shmaps_extra_op -func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { - // CHECK-NOT: call @xla.sdy.manual_computation_body_4 - // CHECK-NOT: call @xla.sdy.manual_computation_body_5 - // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{"a"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"a"} - // CHECK-SAME: (%arg1: tensor<2x8xf32>) { - // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{}, {"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: tensor<2x4xf32>) { - // CHECK-NEXT: %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> - // CHECK-NEXT: sdy.return %[[MUL]] : tensor<2x4xf32> - // CHECK-NEXT: } : (tensor<2x8xf32>) -> tensor<2x8xf32> - // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[MAN_COMP_1]], %[[MAN_COMP_1]] : tensor<2x8xf32> - // CHECK-NEXT: sdy.return %[[ADD]] : tensor<2x8xf32> - // CHECK-NEXT: } : (tensor<4x8xf32>) -> tensor<4x8xf32> - // CHECK-NEXT: return %[[MAN_COMP_0]] : tensor<4x8xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4x8xf32>) -> tensor<2x8xf32> - %1 = call @xla.sdy.manual_computation_body_5(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32> - return %2 : tensor<4x8xf32> -} - -// Tests manual computation with no inputs: there is no `GlobalToLocalShape`, so -// the body call op is the only carrier of the manual axes. -// CHECK-LABEL: func @manual_computation_no_inputs -func.func @manual_computation_no_inputs() -> tensor<4xi64> { - - // CHECK-NOT: call @xla.sdy.manual_computation_body_6 - // CHECK: %[[SHMAP:.*]] = sdy.manual_computation() - // CHECK-SAME{LITERAL}: in_shardings=[] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME{LITERAL}: () { - // CHECK-NEXT: %[[C:.*]] = stablehlo.constant dense<[2, 3]> : tensor<2xi64> - // CHECK-NEXT: sdy.return %[[C]] : tensor<2xi64> - // CHECK-NEXT: } : () -> tensor<4xi64> - // CHECK-NEXT: return %[[SHMAP]] : tensor<4xi64> - %0 = call @xla.sdy.manual_computation_body_6() {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : () -> tensor<2xi64> - %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<4xi64> - return %1 : tensor<4xi64> -} - -// Edge case: Manual computation consuming an input but producing no outputs. -// CHECK-LABEL: func @manual_computation_no_outputs -func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) { - // CHECK-NOT: call @xla.sdy.manual_computation_body_7 - // CHECK: sdy.manual_computation(%arg0) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME{LITERAL}: (%arg1: tensor<2xi64>) { - // CHECK-NEXT: stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> () - // CHECK-NEXT: sdy.return - // CHECK-NEXT: } : (tensor<4xi64>) -> () - // CHECK-NEXT: return - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<4xi64>) -> tensor<4xi64> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4xi64>) -> tensor<2xi64> - call @xla.sdy.manual_computation_body_7(%0) : (tensor<2xi64>) -> () - return -} - -// Edge case: Degenerate manual computation with no inputs and no outputs. -// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs -func.func @manual_computation_no_inputs_no_outputs() { - // CHECK-NEXT: sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () { - // CHECK-NEXT: sdy.return - // CHECK-NEXT: } : () -> () - // CHECK-NEXT: return - call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "false"}} : () -> () - return -} - -// Edge case: Manual computation with 0-sized dimension tensors. -// CHECK-LABEL: func @manual_computation_some_zero_dim_inputs_outputs -func.func @manual_computation_some_zero_dim_inputs_outputs(%arg0: tensor<0x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) { - - // CHECK-NOT: call @xla.sdy.manual_computation_body - // CHECK: %[[CONST_0_32:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32> - // CHECK: %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{}, {"b"}]>, <@mesh_0, [{"b"}, {}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: tensor<0x8xf32>, %arg3: tensor<8x32xf32>) { - // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %arg2, %arg3 - // CHECK-NEXT: sdy.return %[[DOT]], %arg3 - // CHECK-NEXT: } : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) - // CHECK-NEXT: return %[[CONST_0_32]], %[[MAN_COMP]]#1 - %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<0x16xf32>) -> tensor<0x16xf32> - %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32> - %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x8xf32>, tensor<8x32xf32>) - %1:2 = call @xla.sdy.manual_computation_body_9(%c1, %0#1) : (tensor<0x8xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) - %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c2, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>]>} : (tensor<0x32xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) - return %c2, %2#1 : tensor<0x32xf32>, tensor<16x32xf32> -} - -// Type preservation: Token arguments are preserved with empty sharding and not assigned manual axes. -// CHECK-LABEL: func @manual_computation_tokens -func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) { - // CHECK-NOT: call @xla.sdy.manual_computation_body_11 - // CHECK: %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1) - // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] - // CHECK-SAME{LITERAL}: manual_axes={"b"} - // CHECK-SAME: (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) { - // CHECK-NEXT: sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64> - // CHECK-NEXT: } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) - // CHECK-NEXT: return %[[MAN_COMP]]#0, %[[MAN_COMP]]#1 : !stablehlo.token, tensor<2xi64> - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>]>} : (!stablehlo.token) -> !stablehlo.token - %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<2xi64> - %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>) - %1:2 = call @xla.sdy.manual_computation_body_11(%0#0, %0#1) : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) - %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1#0, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{"b"}]>]>} : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>) - return %2#0, %2#1 : !stablehlo.token, tensor<2xi64> -} - -// Tests stray/dangling boundary custom calls where results are unused: safely stripped. -// CHECK-LABEL: func @stray_unused_manual_computation_custom_calls -func.func @stray_unused_manual_computation_custom_calls(%arg0: tensor<0x16xf32>) -> tensor<0x16xf32> { - // CHECK-NEXT: %[[CONST_0_8:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - // CHECK-NEXT: %[[CONST_0_16:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32> - // CHECK-NOT: call @xla.sdy.manual_computation_body - // CHECK-NOT: call @xla.sdy.GlobalToLocalShape - // CHECK-NOT: call @xla.sdy.LocalToGlobalShape - // CHECK-NEXT: return %[[CONST_0_16]] - %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> - %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>) -> tensor<0x8xf32> - %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<0x8xf32>) -> tensor<0x16xf32> - return %c2 : tensor<0x16xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body( -func.func private @xla.sdy.manual_computation_body(%arg0: tensor<2x8xf32>, %arg1: tensor<8x32xf32>) -> tensor<2x32xf32> { - %0 = stablehlo.add %arg0, %arg0 : tensor<2x8xf32> - %1 = stablehlo.dot %0, %arg1 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> - %2 = "stablehlo.all_reduce"(%1) <{replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>}> ({ - ^bb0(%arg2: tensor, %arg3: tensor): - %3 = stablehlo.add %arg2, %arg3 : tensor - stablehlo.return %3 : tensor - }) : (tensor<2x32xf32>) -> tensor<2x32xf32> - return %2 : tensor<2x32xf32> -} - -// CHECK-NOT: func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234( -func.func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - return %arg0 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_0( -func.func private @xla.sdy.manual_computation_body_0(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - return %arg0 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_1( -func.func private @xla.sdy.manual_computation_body_1(%arg0: tensor<8x4xf32>) -> tensor<8x4xf32> { - return %arg0 : tensor<8x4xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_2( -func.func private @xla.sdy.manual_computation_body_2(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { - %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32> - return %0 : tensor<2x4xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_3( -func.func private @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<2x8xf32>) -> tensor<2x4xf32> - %1 = call @xla.sdy.manual_computation_body_2(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32> - return %2 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_4( -func.func private @xla.sdy.manual_computation_body_4(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { - %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32> - return %0 : tensor<2x4xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_5( -func.func private @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { - %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32> - %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<2x8xf32>) -> tensor<2x4xf32> - %1 = call @xla.sdy.manual_computation_body_4(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32> - %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32> - %3 = stablehlo.add %2, %2 : tensor<2x8xf32> - return %3 : tensor<2x8xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_6( -func.func private @xla.sdy.manual_computation_body_6() -> tensor<2xi64> { - %0 = stablehlo.constant dense<[2, 3]> : tensor<2xi64> - return %0 : tensor<2xi64> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_7( -func.func private @xla.sdy.manual_computation_body_7(%arg0: tensor<2xi64>) { - stablehlo.custom_call @sdy_testonly(%arg0) : (tensor<2xi64>) -> () - return -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_8( -func.func private @xla.sdy.manual_computation_body_8() { - return -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_9( -func.func private @xla.sdy.manual_computation_body_9(%arg0: tensor<0x8xf32>, %arg1: tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) { - %0 = stablehlo.dot %arg0, %arg1 : (tensor<0x8xf32>, tensor<8x32xf32>) -> tensor<0x32xf32> - return %0, %arg1 : tensor<0x32xf32>, tensor<8x32xf32> -} - -// CHECK-NOT: func private @xla.sdy.manual_computation_body_11( -func.func private @xla.sdy.manual_computation_body_11(%arg0: !stablehlo.token, %arg1: tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) { - return %arg0, %arg1 : !stablehlo.token, tensor<1xi64> -} From 4ec888cd4f57b9f0a7cb5b7d77d76f61c34ee87d Mon Sep 17 00:00:00 2001 From: Alex Pivovarov Date: Tue, 29 Sep 2026 09:05:57 -0700 Subject: [PATCH 20/25] [GPU] Optimize key conversion in GPU TopK kernels. Store the order-preserving integer key directly in `Descending::KVT` so `Descending::cmp` performs direct integer comparisons without per-comparison bitwise conversions. Apply `details::ToOrdered` once per element when loading keys from global memory in `PerWarpTopK`, and invert the mapping via `details::FromOrdered` only when thread 0 writes the final k outputs in `MergeTopKs`. This reduces the total-order overhead from 1.35x down to 1.12x in comparison to the original C++ less comparator. PiperOrigin-RevId: 990342348 --- .../cuda/topk_kernel_cuda_common.cu.h | 27 ++++++++++++++----- .../rocm/topk_kernel_rocm_common.cu.h | 27 ++++++++++++++----- 2 files changed, 40 insertions(+), 14 deletions(-) diff --git a/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h b/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h index 046e65812860d8..3bb05e6822e6dc 100644 --- a/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h +++ b/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h @@ -81,6 +81,20 @@ __device__ __forceinline__ auto ToOrdered(T x) { return x; } } + +template +__device__ __forceinline__ T FromOrdered(OrderedT val) { + if constexpr (sizeof(T) == 4 && !std::is_integral_v) { + uint32_t u = (val & 0x80000000u) ? (val ^ 0x80000000u) : ~val; + return absl::bit_cast(u); + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + uint16_t u = (val & 0x8000u) ? static_cast(val ^ 0x8000u) + : static_cast(~val); + return absl::bit_cast(u); + } else { + return val; + } +} } // namespace details // Default implementation for KV holder. Useful for testing while adding support @@ -88,15 +102,14 @@ __device__ __forceinline__ auto ToOrdered(T x) { // implementations below. template struct Descending { + using OrderedKey = decltype(details::ToOrdered(T{})); struct KVT { - T key; + OrderedKey key; V idx; }; __device__ __forceinline__ static bool cmp(const KVT& lhs, const KVT& rhs) { - auto l = details::ToOrdered(lhs.key); - auto r = details::ToOrdered(rhs.key); - return l == r ? lhs.idx < rhs.idx : l > r; + return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key; } }; @@ -189,7 +202,7 @@ struct TopK { // TODO(doak): Use bitonic sort. #pragma unroll for (int i = 0; i < K; i++) { - tmp[i] = {key[Idx(i)], VT(Idx(i))}; + tmp[i] = {details::ToOrdered(key[Idx(i)]), VT(Idx(i))}; } #pragma unroll for (int i = 0; i < K; i++) { @@ -205,7 +218,7 @@ struct TopK { constexpr uint32_t WarpSize = WAVEFRONT_SIZE; for (int idx = K; idx < n; idx++) { - KVT kv{key[Idx(idx)], VT(Idx(idx))}; + KVT kv{details::ToOrdered(key[Idx(idx)]), VT(Idx(idx))}; Push(tmp, kv); } Reduce(tmp, WarpSize); @@ -239,7 +252,7 @@ struct TopK { return; } for (int i = 0; i < num_outputs_; ++i) { - keys[i] = tmp[i].key; + keys[i] = details::FromOrdered(tmp[i].key); idxs[i] = tmp[i].idx; } } diff --git a/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h b/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h index df2b8db853fa68..6742958e8d85c8 100644 --- a/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h +++ b/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h @@ -85,6 +85,20 @@ __device__ __forceinline__ auto ToOrdered(T x) { return x; } } + +template +__device__ __forceinline__ T FromOrdered(OrderedT val) { + if constexpr (sizeof(T) == 4 && !std::is_integral_v) { + uint32_t u = (val & 0x80000000u) ? (val ^ 0x80000000u) : ~val; + return absl::bit_cast(u); + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + uint16_t u = (val & 0x8000u) ? static_cast(val ^ 0x8000u) + : static_cast(~val); + return absl::bit_cast(u); + } else { + return val; + } +} } // namespace details // Default implementation for KV holder. Useful for testing while adding support @@ -92,15 +106,14 @@ __device__ __forceinline__ auto ToOrdered(T x) { // implementations below. template struct Descending { + using OrderedKey = decltype(details::ToOrdered(T{})); struct KVT { - T key; + OrderedKey key; V idx; }; __device__ __forceinline__ static bool cmp(const KVT& lhs, const KVT& rhs) { - auto l = details::ToOrdered(lhs.key); - auto r = details::ToOrdered(rhs.key); - return l == r ? lhs.idx < rhs.idx : l > r; + return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key; } }; @@ -193,7 +206,7 @@ struct TopK { // TODO(doak): Use bitonic sort. #pragma unroll for (int i = 0; i < K; i++) { - tmp[i] = {key[Idx(i)], VT(Idx(i))}; + tmp[i] = {details::ToOrdered(key[Idx(i)]), VT(Idx(i))}; } #pragma unroll for (int i = 0; i < K; i++) { @@ -209,7 +222,7 @@ struct TopK { constexpr uint32_t WarpSize = WAVEFRONT_SIZE; for (int idx = K; idx < n; idx++) { - KVT kv{key[Idx(idx)], VT(Idx(idx))}; + KVT kv{details::ToOrdered(key[Idx(idx)]), VT(Idx(idx))}; Push(tmp, kv); } Reduce(tmp, WarpSize); @@ -237,7 +250,7 @@ struct TopK { Reduce(tmp, blockDim.x / WarpSize); if (threadIdx.x != 0) return; for (int i = 0; i < num_outputs_; ++i) { - keys[i] = tmp[i].key; + keys[i] = details::FromOrdered(tmp[i].key); idxs[i] = tmp[i].idx; } } From c3879c337fe4b9accb5f1fb1cb10dd233c2a1830 Mon Sep 17 00:00:00 2001 From: "Patrick C. Toulme" Date: Tue, 29 Sep 2026 09:20:56 -0700 Subject: [PATCH 21/25] [XLA] Track dead while parameters with dense slot arrays and bit sets TryRemoveDeadWhileParams kept a bit set of all tuple indices and a UnionFind entry per instruction in two hash maps and ORed the bit sets along every operand edge. Dense arrays indexed by the instruction's local id replace the maps, the bit sets cover only the removable candidate indices, and one post order pass per computation fills them. Same removed indices, same HLO. PiperOrigin-RevId: 990350581 --- third_party/xla/xla/BUILD | 23 +- third_party/xla/xla/bit_set.h | 180 +++++++ third_party/xla/xla/bit_set_test.cc | 404 ++++++++++++++++ third_party/xla/xla/inlined_bit_set.h | 95 ---- third_party/xla/xla/inlined_bit_set_test.cc | 196 -------- third_party/xla/xla/service/BUILD | 3 +- .../xla/xla/service/while_loop_simplifier.cc | 444 ++++++++++++------ .../xla/service/while_loop_simplifier_test.cc | 69 +++ third_party/xla/xla/union_find.h | 75 +++ third_party/xla/xla/union_find_test.cc | 209 +++++++++ 10 files changed, 1253 insertions(+), 445 deletions(-) create mode 100644 third_party/xla/xla/bit_set.h create mode 100644 third_party/xla/xla/bit_set_test.cc delete mode 100644 third_party/xla/xla/inlined_bit_set.h delete mode 100644 third_party/xla/xla/inlined_bit_set_test.cc create mode 100644 third_party/xla/xla/union_find_test.cc diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 406a441dc69dcc..3da1bb89bd9bc2 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1425,11 +1425,23 @@ xla_cc_test( cc_library( name = "union_find", hdrs = ["union_find.h"], + deps = ["@com_google_absl//absl/log:check"], +) + +xla_cc_test( + name = "union_find_test", + srcs = ["union_find_test.cc"], + deps = [ + ":union_find", + "//xla/tsl/platform:test", + "//xla/tsl/platform:test_benchmark", + "//xla/tsl/platform:test_main", + ], ) cc_library( - name = "inlined_bit_set", - hdrs = ["inlined_bit_set.h"], + name = "bit_set", + hdrs = ["bit_set.h"], deps = [ "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log:check", @@ -1437,11 +1449,12 @@ cc_library( ) xla_cc_test( - name = "inlined_bit_set_test", - srcs = ["inlined_bit_set_test.cc"], + name = "bit_set_test", + srcs = ["bit_set_test.cc"], deps = [ - ":inlined_bit_set", + ":bit_set", "//xla/tsl/platform:test", + "//xla/tsl/platform:test_benchmark", "//xla/tsl/platform:test_main", "@com_google_absl//absl/container:flat_hash_set", ], diff --git a/third_party/xla/xla/bit_set.h b/third_party/xla/xla/bit_set.h new file mode 100644 index 00000000000000..875d5cabc22a27 --- /dev/null +++ b/third_party/xla/xla/bit_set.h @@ -0,0 +1,180 @@ +/* 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_BIT_SET_H_ +#define XLA_BIT_SET_H_ + +#include +#include +#include +#include + +#include "absl/container/inlined_vector.h" +#include "absl/log/check.h" + +namespace xla { + +// A fixed-size bit set (or bitmap) that stores up to `InlinedWords * 64` bits +// inline without heap allocation, spilling to the heap only when `num_bits` +// exceeds this capacity. Useful for tracking sets of dense integer indices +// over a typically small range `[0, num_bits)`. +template +class InlinedBitSet { + public: + using Word = uint64_t; + static constexpr int kWordSize = sizeof(Word) * 8; + + InlinedBitSet() = default; + explicit InlinedBitSet(int num_bits) : words_(NumWords(num_bits), 0) {} + + void Set(int index) { + DCHECK_GE(index, 0); + const size_t word_idx = IndexToWord(index); + words_[word_idx] |= (Word{1} << (index % kWordSize)); + } + + // Sets all bits in [0, num_bits) to 1 + void SetAll(int num_bits) { + DCHECK_GE(num_bits, 0); + DCHECK_LE(static_cast(num_bits), words_.size() * kWordSize); + const size_t full_words = num_bits / kWordSize; + std::fill_n(words_.begin(), full_words, ~Word{0}); + const size_t remainder = num_bits % kWordSize; + if (remainder != 0) { + words_[full_words] |= (Word{1} << remainder) - 1; + } + } + + void Clear(int index) { + const size_t word_idx = IndexToWord(index); + words_[word_idx] &= ~(Word{1} << (index % kWordSize)); + } + + bool Test(int index) const { + const size_t word_idx = IndexToWord(index); + return (words_[word_idx] & (Word{1} << (index % kWordSize))) != 0; + } + + InlinedBitSet& operator|=(const InlinedBitSet& other) { + DCHECK_EQ(words_.size(), other.words_.size()); + for (size_t i = 0; i < words_.size(); ++i) { + words_[i] |= other.words_[i]; + } + return *this; + } + + private: + static size_t NumWords(int num_bits) { + DCHECK_GE(num_bits, 0); + return (num_bits + kWordSize - 1) / kWordSize; + } + + size_t IndexToWord(int index) const { + DCHECK_GE(index, 0); + const size_t word_idx = index / kWordSize; + DCHECK_LT(word_idx, words_.size()); + return word_idx; + } + + absl::InlinedVector words_; +}; + +// Holds num_rows bit sets of num_bits bits each in a single heap allocation. +// Useful when an analysis keeps one bit set (of the same size) per element of +// a large dense index space, for example per instruction: one allocation +// instead of one per element, and the rows sit next to each other in memory, +// so an OR of one row into another streams over two contiguous word ranges. +// Rows are addressed by index. +// +// Supported ops: +// - Set: sets one bit of a row. +// - Test: reads one bit of a row. +// - Or: ORs one row into another. +// - OrIgnoringBit: ORs one row into another, leaving one bit out. +class DenseBitSets { + public: + using Word = uint64_t; + static constexpr int kWordSize = sizeof(Word) * 8; + + DenseBitSets(int num_rows, int num_bits) + : num_rows_(CheckNonNegative(num_rows)), + num_bits_(CheckNonNegative(num_bits)), + words_per_row_(num_bits_ / kWordSize + (num_bits_ % kWordSize != 0)), + words_(static_cast(num_rows_) * words_per_row_, 0) {} + + int num_rows() const { return num_rows_; } + int num_bits() const { return num_bits_; } + + // Sets bit 'bit' of row 'row'. + void Set(int row, int bit) { Row(row)[WordIndex(bit)] |= Mask(bit); } + + // Returns whether bit 'bit' of row 'row' is set. + bool Test(int row, int bit) const { + return (Row(row)[WordIndex(bit)] & Mask(bit)) != 0; + } + + // row |= other. + void Or(int row, int other) { + Word* dst = Row(row); + const Word* src = Row(other); + for (int w = 0; w < words_per_row_; ++w) { + dst[w] |= src[w]; + } + } + + // row |= other, with ignored_bit of other left out: a plain Or, which the + // compiler vectorizes, then the destination's own value of the bit is put + // back. + void OrIgnoringBit(int row, int other, int ignored_bit) { + Word& word = Row(row)[WordIndex(ignored_bit)]; + const Word own_bit = word & Mask(ignored_bit); + Or(row, other); + word = (word & ~Mask(ignored_bit)) | own_bit; + } + + private: + static int CheckNonNegative(int count) { + DCHECK_GE(count, 0); + return count; + } + + static constexpr Word Mask(int bit) { return Word{1} << (bit % kWordSize); } + + int WordIndex(int bit) const { + DCHECK_GE(bit, 0); + DCHECK_LT(bit, num_bits_); + return bit / kWordSize; + } + + // The word offset of a row; computed in size_t because the whole block can + // hold more than 2^31 words. + size_t RowOffset(int row) const { + DCHECK_GE(row, 0); + DCHECK_LT(row, num_rows_); + return static_cast(row) * words_per_row_; + } + + Word* Row(int row) { return words_.data() + RowOffset(row); } + const Word* Row(int row) const { return words_.data() + RowOffset(row); } + + int num_rows_; + int num_bits_; + int words_per_row_; + std::vector words_; +}; + +} // namespace xla + +#endif // XLA_BIT_SET_H_ diff --git a/third_party/xla/xla/bit_set_test.cc b/third_party/xla/xla/bit_set_test.cc new file mode 100644 index 00000000000000..e9a5a634e034bf --- /dev/null +++ b/third_party/xla/xla/bit_set_test.cc @@ -0,0 +1,404 @@ +/* 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/bit_set.h" + +#include +#include + +#include "absl/container/flat_hash_set.h" +#include "xla/tsl/platform/test.h" +#include "xla/tsl/platform/test_benchmark.h" + +namespace xla { +namespace { + +using ::benchmark::DoNotOptimize; +using ::testing::benchmark::State; + +class BruteForceBitSet { + public: + void Set(int index) { bits_.insert(index); } + + void SetAll(int num_bits) { + for (int i = 0; i < num_bits; ++i) { + bits_.insert(i); + } + } + + void Clear(int index) { bits_.erase(index); } + + bool Test(int index) const { return bits_.contains(index); } + + BruteForceBitSet& operator|=(const BruteForceBitSet& other) { + bits_.insert(other.bits_.begin(), other.bits_.end()); + return *this; + } + + private: + absl::flat_hash_set bits_; +}; + +TEST(InlinedBitSetTest, DefaultInitializedAllFalse) { + InlinedBitSet<> bits(130); + + EXPECT_FALSE(bits.Test(0)); + EXPECT_FALSE(bits.Test(63)); + EXPECT_FALSE(bits.Test(64)); + EXPECT_FALSE(bits.Test(127)); + EXPECT_FALSE(bits.Test(128)); +} + +TEST(InlinedBitSetTest, Set) { + InlinedBitSet<> bits(130); + + bits.Set(0); + bits.Set(63); + bits.Set(64); + bits.Set(128); + + EXPECT_TRUE(bits.Test(0)); + EXPECT_TRUE(bits.Test(63)); + EXPECT_TRUE(bits.Test(64)); + EXPECT_FALSE(bits.Test(65)); + EXPECT_FALSE(bits.Test(127)); + EXPECT_TRUE(bits.Test(128)); +} + +TEST(InlinedBitSetTest, Clear) { + InlinedBitSet<> bits(130); + bits.Set(63); + bits.Set(64); + + bits.Clear(63); + + EXPECT_FALSE(bits.Test(63)); + EXPECT_TRUE(bits.Test(64)); +} + +TEST(InlinedBitSetTest, SetAllZeroBits) { + InlinedBitSet<> bits(0); + + bits.SetAll(0); +} + +TEST(InlinedBitSetTest, SetAllExactWordBoundary) { + InlinedBitSet<> bits(64); + + bits.SetAll(64); + + for (int i = 0; i < 64; ++i) { + EXPECT_TRUE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, SetAllNonExactBoundary) { + InlinedBitSet<> bits(70); + + bits.SetAll(70); + + for (int i = 0; i < 70; ++i) { + EXPECT_TRUE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, SetAllLessThanCapacity) { + InlinedBitSet<> bits(128); + + bits.SetAll(50); + + for (int i = 0; i < 50; ++i) { + EXPECT_TRUE(bits.Test(i)); + } + for (int i = 50; i < 128; ++i) { + EXPECT_FALSE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, SetAllZeroOnNonEmpty) { + InlinedBitSet<> bits(64); + + bits.SetAll(0); + + for (int i = 0; i < 64; ++i) { + EXPECT_FALSE(bits.Test(i)); + } +} + +TEST(InlinedBitSetTest, BitwiseOr) { + InlinedBitSet<> a(130); + InlinedBitSet<> b(130); + a.Set(10); + a.Set(100); + b.Set(70); + b.Set(100); + + a |= b; + + EXPECT_TRUE(a.Test(10)); + EXPECT_TRUE(a.Test(70)); + EXPECT_TRUE(a.Test(100)); + EXPECT_FALSE(a.Test(0)); + EXPECT_FALSE(a.Test(129)); +} + +TEST(InlinedBitSetTest, MatchesBruteForceOnRandomOperations) { + constexpr int kNumBits = 200; + constexpr int kNumOps = 500; + std::mt19937 rng(42); + std::uniform_int_distribution op_dist(0, 3); + std::uniform_int_distribution bit_dist(0, kNumBits - 1); + std::uniform_int_distribution count_dist(0, kNumBits); + InlinedBitSet<> bits(kNumBits); + InlinedBitSet<> other_bits(kNumBits); + BruteForceBitSet bf_bits; + BruteForceBitSet bf_other_bits; + + for (int step = 0; step < kNumOps; ++step) { + const int idx = bit_dist(rng); + other_bits.Set(idx); + bf_other_bits.Set(idx); + switch (op_dist(rng)) { + case 0: + bits.Set(idx); + bf_bits.Set(idx); + break; + case 1: + bits.Clear(idx); + bf_bits.Clear(idx); + break; + case 2: { + const int count = count_dist(rng); + bits.SetAll(count); + bf_bits.SetAll(count); + break; + } + case 3: + bits |= other_bits; + bf_bits |= bf_other_bits; + break; + } + } + + for (int i = 0; i < kNumBits; ++i) { + EXPECT_EQ(bits.Test(i), bf_bits.Test(i)) << "Mismatch at index " << i; + } +} + +// Reference implementation of DenseBitSets: rows of bits as a vector of +// vectors. +class BruteForceBitSetRows { + public: + BruteForceBitSetRows(int num_rows, int num_bits) + : rows_(num_rows, std::vector(num_bits, false)) {} + + void Set(int row, int bit) { rows_[row][bit] = true; } + bool Test(int row, int bit) const { return rows_[row][bit]; } + + void Or(int row, int other) { + for (int bit = 0; bit < rows_[row].size(); ++bit) { + rows_[row][bit] = rows_[row][bit] || rows_[other][bit]; + } + } + + void OrIgnoringBit(int row, int other, int ignored_bit) { + for (int bit = 0; bit < rows_[row].size(); ++bit) { + if (bit != ignored_bit) { + rows_[row][bit] = rows_[row][bit] || rows_[other][bit]; + } + } + } + + private: + std::vector> rows_; +}; + +TEST(DenseBitSetsTest, EmptyShapesAreAllowed) { + DenseBitSets no_rows(0, 130); + EXPECT_EQ(no_rows.num_rows(), 0); + EXPECT_EQ(no_rows.num_bits(), 130); + + DenseBitSets no_bits(3, 0); + EXPECT_EQ(no_bits.num_rows(), 3); + EXPECT_EQ(no_bits.num_bits(), 0); + no_bits.Or(0, 1); // Rows of zero words. + + DenseBitSets empty(0, 0); + EXPECT_EQ(empty.num_rows(), 0); + EXPECT_EQ(empty.num_bits(), 0); +} + +TEST(DenseBitSetsTest, DefaultInitializedAllFalse) { + DenseBitSets bits(3, 130); + + EXPECT_EQ(bits.num_rows(), 3); + EXPECT_EQ(bits.num_bits(), 130); + for (int row = 0; row < 3; ++row) { + EXPECT_FALSE(bits.Test(row, 0)); + EXPECT_FALSE(bits.Test(row, 63)); + EXPECT_FALSE(bits.Test(row, 64)); + EXPECT_FALSE(bits.Test(row, 129)); + } +} + +TEST(DenseBitSetsTest, SetIsConfinedToItsRow) { + DenseBitSets bits(3, 130); + + bits.Set(1, 0); + bits.Set(1, 63); + bits.Set(1, 64); + bits.Set(1, 129); + + EXPECT_TRUE(bits.Test(1, 0)); + EXPECT_TRUE(bits.Test(1, 63)); + EXPECT_TRUE(bits.Test(1, 64)); + EXPECT_TRUE(bits.Test(1, 129)); + EXPECT_FALSE(bits.Test(1, 1)); + EXPECT_FALSE(bits.Test(1, 128)); + for (int bit : {0, 63, 64, 129}) { + EXPECT_FALSE(bits.Test(0, bit)); + EXPECT_FALSE(bits.Test(2, bit)); + } +} + +TEST(DenseBitSetsTest, OrAndOrIgnoringBit) { + DenseBitSets bits(2, 130); + bits.Set(0, 5); + bits.Set(1, 5); + bits.Set(1, 64); + bits.Set(1, 129); + + bits.Or(0, 1); + + EXPECT_TRUE(bits.Test(0, 5)); + EXPECT_TRUE(bits.Test(0, 64)); + EXPECT_TRUE(bits.Test(0, 129)); + EXPECT_FALSE(bits.Test(0, 6)); + + DenseBitSets more(2, 130); + more.Set(0, 64); + more.Set(1, 64); + more.Set(1, 65); + more.Set(1, 129); + + more.OrIgnoringBit(0, 1, 129); + more.OrIgnoringBit(0, 1, 64); + + EXPECT_TRUE(more.Test(0, 64)); // Already set in the row; never cleared. + EXPECT_TRUE(more.Test(0, 65)); + EXPECT_TRUE(more.Test(0, 129)); // Copied by the second call. + EXPECT_TRUE(more.Test(1, 64)); +} + +TEST(DenseBitSetsTest, MatchesBruteForceOnRandomOperations) { + constexpr int kNumInstances = 200; + constexpr int kNumOps = 300; + std::mt19937 rng(7); + std::uniform_int_distribution rows_dist(1, 12); + std::uniform_int_distribution bits_dist(1, 300); + std::uniform_int_distribution op_dist(0, 2); + + for (int instance = 0; instance < kNumInstances; ++instance) { + const int num_rows = rows_dist(rng); + // Every fourth instance sits on a word boundary. + const int num_bits = + instance % 4 == 0 ? 64 * (1 + instance % 3) : bits_dist(rng); + std::uniform_int_distribution row_dist(0, num_rows - 1); + std::uniform_int_distribution bit_dist(0, num_bits - 1); + DenseBitSets bits(num_rows, num_bits); + BruteForceBitSetRows bf_bits(num_rows, num_bits); + + for (int step = 0; step < kNumOps; ++step) { + const int row = row_dist(rng); + const int other = row_dist(rng); + const int bit = bit_dist(rng); + switch (op_dist(rng)) { + case 0: + bits.Set(row, bit); + bf_bits.Set(row, bit); + break; + case 1: + bits.Or(row, other); + bf_bits.Or(row, other); + break; + case 2: + bits.OrIgnoringBit(row, other, bit); + bf_bits.OrIgnoringBit(row, other, bit); + break; + } + ASSERT_EQ(bits.Test(row, bit), bf_bits.Test(row, bit)) + << "instance " << instance << " step " << step << " row " << row + << " bit " << bit; + } + for (int row = 0; row < num_rows; ++row) { + for (int bit = 0; bit < num_bits; ++bit) { + ASSERT_EQ(bits.Test(row, bit), bf_bits.Test(row, bit)) + << "instance " << instance << " row " << row << " bit " << bit; + } + } + } +} + +// Benchmarks: the operations TryRemoveDeadWhileParams issues, for rows of 64 +// to 65,536 bits (more than twice the largest while tuple seen so far): +// - Or, once per operand edge, +// - OrIgnoringBit, once per body output, +// - Set and Test, once per get-tuple-element. +void BM_DenseBitSetsOr(State& state) { + const int num_bits = state.range(0); + DenseBitSets bits(2, num_bits); + bits.Set(1, num_bits - 1); + for (auto s : state) { + bits.Or(0, 1); + DoNotOptimize(bits); + } + state.SetBytesProcessed(state.iterations() * ((num_bits + 63) / 64) * 8); +} +BENCHMARK(BM_DenseBitSetsOr)->Arg(64)->Arg(1024)->Arg(16384)->Arg(65536); + +void BM_DenseBitSetsOrIgnoringBit(State& state) { + const int num_bits = state.range(0); + DenseBitSets bits(2, num_bits); + bits.Set(1, num_bits - 1); + for (auto s : state) { + bits.OrIgnoringBit(0, 1, num_bits / 2); + DoNotOptimize(bits); + } + state.SetBytesProcessed(state.iterations() * ((num_bits + 63) / 64) * 8); +} +BENCHMARK(BM_DenseBitSetsOrIgnoringBit) + ->Arg(64) + ->Arg(1024) + ->Arg(16384) + ->Arg(65536); + +void BM_DenseBitSetsSetTest(State& state) { + const int num_bits = state.range(0); + DenseBitSets bits(1024, num_bits); + int row = 0; + int bit = 0; + bool sink = false; + for (auto s : state) { + bits.Set(row, bit); + sink ^= bits.Test((row + 1) % 1024, bit); + row = (row + 7) % 1024; + bit = (bit + 13) % num_bits; + } + DoNotOptimize(sink); +} +BENCHMARK(BM_DenseBitSetsSetTest)->Arg(64)->Arg(1024)->Arg(16384)->Arg(65536); + +} // namespace +} // namespace xla diff --git a/third_party/xla/xla/inlined_bit_set.h b/third_party/xla/xla/inlined_bit_set.h deleted file mode 100644 index 468c7c692708c5..00000000000000 --- a/third_party/xla/xla/inlined_bit_set.h +++ /dev/null @@ -1,95 +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_INLINED_BIT_SET_H_ -#define XLA_INLINED_BIT_SET_H_ - -#include -#include -#include - -#include "absl/container/inlined_vector.h" -#include "absl/log/check.h" - -namespace xla { - -// A fixed-size bit set (or bitmap) that stores up to `InlinedWords * 64` bits -// inline without heap allocation, spilling to the heap only when `num_bits` -// exceeds this capacity. Useful for tracking sets of dense integer indices -// over a typically small range `[0, num_bits)`. -template -class InlinedBitSet { - public: - using Word = uint64_t; - static constexpr int kWordSize = sizeof(Word) * 8; - - InlinedBitSet() = default; - explicit InlinedBitSet(int num_bits) : words_(NumWords(num_bits), 0) {} - - void Set(int index) { - DCHECK_GE(index, 0); - const size_t word_idx = IndexToWord(index); - words_[word_idx] |= (Word{1} << (index % kWordSize)); - } - - // Sets all bits in [0, num_bits) to 1 - void SetAll(int num_bits) { - DCHECK_GE(num_bits, 0); - DCHECK_LE(static_cast(num_bits), words_.size() * kWordSize); - const size_t full_words = num_bits / kWordSize; - std::fill_n(words_.begin(), full_words, ~Word{0}); - const size_t remainder = num_bits % kWordSize; - if (remainder != 0) { - words_[full_words] |= (Word{1} << remainder) - 1; - } - } - - void Clear(int index) { - const size_t word_idx = IndexToWord(index); - words_[word_idx] &= ~(Word{1} << (index % kWordSize)); - } - - bool Test(int index) const { - const size_t word_idx = IndexToWord(index); - return (words_[word_idx] & (Word{1} << (index % kWordSize))) != 0; - } - - InlinedBitSet& operator|=(const InlinedBitSet& other) { - DCHECK_EQ(words_.size(), other.words_.size()); - for (size_t i = 0; i < words_.size(); ++i) { - words_[i] |= other.words_[i]; - } - return *this; - } - - private: - static size_t NumWords(int num_bits) { - DCHECK_GE(num_bits, 0); - return (num_bits + kWordSize - 1) / kWordSize; - } - - size_t IndexToWord(int index) const { - DCHECK_GE(index, 0); - const size_t word_idx = index / kWordSize; - DCHECK_LT(word_idx, words_.size()); - return word_idx; - } - - absl::InlinedVector words_; -}; - -} // namespace xla - -#endif // XLA_INLINED_BIT_SET_H_ diff --git a/third_party/xla/xla/inlined_bit_set_test.cc b/third_party/xla/xla/inlined_bit_set_test.cc deleted file mode 100644 index d80413b5320886..00000000000000 --- a/third_party/xla/xla/inlined_bit_set_test.cc +++ /dev/null @@ -1,196 +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/inlined_bit_set.h" - -#include - -#include "absl/container/flat_hash_set.h" -#include "xla/tsl/platform/test.h" - -namespace xla { -namespace { - -class BruteForceBitSet { - public: - void Set(int index) { bits_.insert(index); } - - void SetAll(int num_bits) { - for (int i = 0; i < num_bits; ++i) { - bits_.insert(i); - } - } - - void Clear(int index) { bits_.erase(index); } - - bool Test(int index) const { return bits_.contains(index); } - - BruteForceBitSet& operator|=(const BruteForceBitSet& other) { - bits_.insert(other.bits_.begin(), other.bits_.end()); - return *this; - } - - private: - absl::flat_hash_set bits_; -}; - -TEST(InlinedBitSetTest, DefaultInitializedAllFalse) { - InlinedBitSet<> bits(130); - - EXPECT_FALSE(bits.Test(0)); - EXPECT_FALSE(bits.Test(63)); - EXPECT_FALSE(bits.Test(64)); - EXPECT_FALSE(bits.Test(127)); - EXPECT_FALSE(bits.Test(128)); -} - -TEST(InlinedBitSetTest, Set) { - InlinedBitSet<> bits(130); - - bits.Set(0); - bits.Set(63); - bits.Set(64); - bits.Set(128); - - EXPECT_TRUE(bits.Test(0)); - EXPECT_TRUE(bits.Test(63)); - EXPECT_TRUE(bits.Test(64)); - EXPECT_FALSE(bits.Test(65)); - EXPECT_FALSE(bits.Test(127)); - EXPECT_TRUE(bits.Test(128)); -} - -TEST(InlinedBitSetTest, Clear) { - InlinedBitSet<> bits(130); - bits.Set(63); - bits.Set(64); - - bits.Clear(63); - - EXPECT_FALSE(bits.Test(63)); - EXPECT_TRUE(bits.Test(64)); -} - -TEST(InlinedBitSetTest, SetAllZeroBits) { - InlinedBitSet<> bits(0); - - bits.SetAll(0); -} - -TEST(InlinedBitSetTest, SetAllExactWordBoundary) { - InlinedBitSet<> bits(64); - - bits.SetAll(64); - - for (int i = 0; i < 64; ++i) { - EXPECT_TRUE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, SetAllNonExactBoundary) { - InlinedBitSet<> bits(70); - - bits.SetAll(70); - - for (int i = 0; i < 70; ++i) { - EXPECT_TRUE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, SetAllLessThanCapacity) { - InlinedBitSet<> bits(128); - - bits.SetAll(50); - - for (int i = 0; i < 50; ++i) { - EXPECT_TRUE(bits.Test(i)); - } - for (int i = 50; i < 128; ++i) { - EXPECT_FALSE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, SetAllZeroOnNonEmpty) { - InlinedBitSet<> bits(64); - - bits.SetAll(0); - - for (int i = 0; i < 64; ++i) { - EXPECT_FALSE(bits.Test(i)); - } -} - -TEST(InlinedBitSetTest, BitwiseOr) { - InlinedBitSet<> a(130); - InlinedBitSet<> b(130); - a.Set(10); - a.Set(100); - b.Set(70); - b.Set(100); - - a |= b; - - EXPECT_TRUE(a.Test(10)); - EXPECT_TRUE(a.Test(70)); - EXPECT_TRUE(a.Test(100)); - EXPECT_FALSE(a.Test(0)); - EXPECT_FALSE(a.Test(129)); -} - -TEST(InlinedBitSetTest, MatchesBruteForceOnRandomOperations) { - constexpr int kNumBits = 200; - constexpr int kNumOps = 500; - std::mt19937 rng(42); - std::uniform_int_distribution op_dist(0, 3); - std::uniform_int_distribution bit_dist(0, kNumBits - 1); - std::uniform_int_distribution count_dist(0, kNumBits); - InlinedBitSet<> bits(kNumBits); - InlinedBitSet<> other_bits(kNumBits); - BruteForceBitSet bf_bits; - BruteForceBitSet bf_other_bits; - - for (int step = 0; step < kNumOps; ++step) { - const int idx = bit_dist(rng); - other_bits.Set(idx); - bf_other_bits.Set(idx); - switch (op_dist(rng)) { - case 0: - bits.Set(idx); - bf_bits.Set(idx); - break; - case 1: - bits.Clear(idx); - bf_bits.Clear(idx); - break; - case 2: { - const int count = count_dist(rng); - bits.SetAll(count); - bf_bits.SetAll(count); - break; - } - case 3: - bits |= other_bits; - bf_bits |= bf_other_bits; - break; - } - } - - for (int i = 0; i < kNumBits; ++i) { - EXPECT_EQ(bits.Test(i), bf_bits.Test(i)) << "Mismatch at index " << i; - } -} - -} // namespace -} // namespace xla diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 0d21749f924035..81299291f0dcc5 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2816,9 +2816,9 @@ cc_library( ":hlo_creation_utils", ":pattern_matcher", ":while_util", + "//xla:bit_set", "//xla:comparison_util", "//xla:frontend_attributes", - "//xla:inlined_bit_set", "//xla:literal_util", "//xla:shape_util", "//xla:status_macros", @@ -2863,6 +2863,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/service/while_loop_simplifier.cc b/third_party/xla/xla/service/while_loop_simplifier.cc index aa1cb7b60455cb..3e99ac8dcbc57d 100644 --- a/third_party/xla/xla/service/while_loop_simplifier.cc +++ b/third_party/xla/xla/service/while_loop_simplifier.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "xla/bit_set.h" #include "xla/comparison_util.h" #include "xla/frontend_attributes.h" #include "xla/hlo/analysis/while_loop_analysis.h" @@ -43,7 +44,6 @@ limitations under the License. #include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/hlo/utils/hlo_query.h" -#include "xla/inlined_bit_set.h" #include "xla/literal_util.h" #include "xla/primitive_util.h" #include "xla/service/call_inliner.h" @@ -143,19 +143,15 @@ void CopyMetadata(HloInstruction* old_while_op, HloInstruction* new_while_op) { new_while_op->set_metadata(old_while_op->metadata()); } -// This is a utility function that removes the given tuple indices from the -// while loop init, body, and condition. The final shape returned is still the -// same as before. If set index_for_replaced will replace any use of the removed -// indices in the final shape with a copy of the removed index. +// This is a utility function that removes the tuple indices whose bit is not +// set in used_tuple_indices from the while loop init, body, and condition. The +// final shape returned is still the same as before. If set index_for_replaced +// will replace any use of the removed indices in the final shape with a copy of +// the removed index. static absl::StatusOr RemoveDeadTupleIndices( - HloInstruction* while_op, absl::flat_hash_set& used_tuple_indices, + HloInstruction* while_op, const InlinedBitSet<>& used_tuple_indices, std::optional> dead_to_surviving_index = std::nullopt) { - // Build up maps from the old/new to the new/old tuple indices. - std::vector new_to_old_tuple_idx(used_tuple_indices.begin(), - used_tuple_indices.end()); - absl::c_sort(new_to_old_tuple_idx); - HloModule* module = while_op->GetModule(); HloComputation* computation = while_op->parent(); HloInstruction* while_init = while_op->mutable_operand(0); @@ -163,6 +159,15 @@ static absl::StatusOr RemoveDeadTupleIndices( HloComputation* while_body = while_op->while_body(); HloInstruction* while_body_root = while_body->root_instruction(); + // Build up maps from the old/new to the new/old tuple indices. + std::vector new_to_old_tuple_idx; + const int64_t old_tuple_size = while_init->shape().tuple_shapes().size(); + for (int64_t old_idx = 0; old_idx < old_tuple_size; ++old_idx) { + if (used_tuple_indices.Test(old_idx)) { + new_to_old_tuple_idx.push_back(old_idx); + } + } + auto print_no_metadata = HloPrintOptions().set_print_metadata(false); absl::flat_hash_map old_to_new_tuple_idx; @@ -355,6 +360,253 @@ bool AllWhileParamConsumersGte(const HloInstruction* while_op) { return true; } +// The tuple indices of a while loop that are read after the loop: all of them +// when the loop is the root of its computation or has a user other than +// get-tuple-element. +InlinedBitSet<> TupleIndicesUsedAfterLoop(const HloInstruction* while_op, + int64_t tuple_size) { + InlinedBitSet<> used(tuple_size); + if (while_op == while_op->parent()->root_instruction()) { + used.SetAll(tuple_size); + return used; + } + for (const HloInstruction* user : while_op->users()) { + if (user->opcode() != HloOpcode::kGetTupleElement) { + used.SetAll(tuple_size); + return used; + } + used.Set(user->tuple_index()); + } + return used; +} + +// Whether the body passes the tuple element at 'index' of its parameter +// through unmodified. +bool IsPassedThroughUnmodified(const HloComputation* while_body, + int64_t index) { + const HloInstruction* output = while_body->root_instruction()->operand(index); + return output->opcode() == HloOpcode::kGetTupleElement && + output->operand(0) == while_body->parameter_instruction(0) && + output->tuple_index() == index; +} + +// How a while loop's body and condition depend on the loop's candidate tuple +// indices, the indices TryRemoveDeadWhileParams may remove: candidate_bit +// numbers them and holds -1 for every other index. Per instruction, the +// candidates it depends on within one iteration; across instructions, the +// connected components of the operand graph. +// +// The instructions of the body and the condition are addressed by slot, the +// body's local ids first, then the condition's. The slots map the instructions +// to a dense index range, so the per instruction state lives in flat arrays +// instead of hash maps keyed by instruction. Local ids stay stable during the +// analysis: nothing mutates the graph before RemoveDeadTupleIndices. +class WhileInputDependencies { + public: + WhileInputDependencies(const HloInstruction* while_op, + absl::Span candidate_bit, + int num_candidates) + : while_body_(while_op->while_body()), + while_cond_(while_op->while_condition()), + while_body_root_(while_body_->root_instruction()), + candidate_bit_(candidate_bit), + body_slots_(while_body_->next_unique_instruction_internal_id()), + num_slots_(body_slots_ + + while_cond_->next_unique_instruction_internal_id()), + side_effecting_row_(num_slots_), + affecting_others_row_(num_slots_ + 1), + input_deps_(num_slots_ + 2, num_candidates), + sets_(num_slots_) { + PropagateThrough(while_body_); + PropagateThrough(while_cond_); + AccumulateOutputDependencies(); + } + + int64_t tuple_size() const { return candidate_bit_.size(); } + + // Whether the index may be removed at all. + bool IsCandidate(int64_t index) const { return candidate_bit_[index] >= 0; } + + // Whether a side effecting instruction or the loop condition depends on the + // candidate index. + bool ReachesSink(int64_t index) const { + DCHECK(IsCandidate(index)); + return input_deps_.Test(side_effecting_row_, candidate_bit_[index]); + } + + // Whether a body output other than its own depends on the candidate index. + bool AffectsOtherOutputs(int64_t index) const { + DCHECK(IsCandidate(index)); + return input_deps_.Test(affecting_others_row_, candidate_bit_[index]); + } + + // The connected component of the body output at the index, numbered below + // num_slots(). The get-tuple-elements of the index belong to the same + // component, so indices that depend on each other share it. + int ComponentOf(int64_t index) { + return sets_.Find(SlotOf(while_body_root_->operand(index))); + } + + int num_slots() const { return num_slots_; } + + private: + int SlotOf(const HloInstruction* inst) const { + DCHECK(inst->parent() == while_body_ || inst->parent() == while_cond_); + return (inst->parent() == while_body_ ? 0 : body_slots_) + inst->local_id(); + } + + // The parameters and the body root take no part in the dependencies. + bool IsTracked(const HloInstruction* inst) const { + return inst != while_body_->parameter_instruction(0) && + inst != while_cond_->parameter_instruction(0) && + inst != while_body_root_; + } + + // One post order pass over a computation. Operands come before their users, + // so the dependencies of an instruction are the union of its operands' by + // the time it is visited. A get-tuple-element of the parameter seeds its own + // bit and joins the component of the output it feeds back into. A side + // effecting instruction and the condition root are sinks: the inputs they + // depend on must stay. + void PropagateThrough(const HloComputation* comp) { + const HloInstruction* while_input = comp->parameter_instruction(0); + comp->ForEachInstructionPostOrder([&](HloInstruction* inst) { + if (!IsTracked(inst)) { + return; + } + const int slot = SlotOf(inst); + if (inst->opcode() == HloOpcode::kGetTupleElement && + inst->operand(0) == while_input) { + if (candidate_bit_[inst->tuple_index()] >= 0) { + input_deps_.Set(slot, candidate_bit_[inst->tuple_index()]); + } + const HloInstruction* output = + while_body_root_->operand(inst->tuple_index()); + if (output != inst) { + sets_.Merge(SlotOf(output), slot); + } + } else { + for (const HloInstruction* operand : inst->operands()) { + if (!IsTracked(operand)) { + continue; + } + sets_.Merge(SlotOf(operand), slot); + if (const std::optional index = + ParameterIndex(operand, while_input)) { + // The row of a parameter get-tuple-element holds only its own + // bit, so one Set replaces the OR of a whole row. + if (candidate_bit_[*index] >= 0) { + input_deps_.Set(slot, candidate_bit_[*index]); + } + } else { + input_deps_.Or(slot, SlotOf(operand)); + } + } + } + if (inst->HasSideEffect() || inst == while_cond_->root_instruction()) { + input_deps_.Or(side_effecting_row_, slot); + } + }); + } + + // The tuple index when inst is a get-tuple-element of the given parameter. + static std::optional ParameterIndex(const HloInstruction* inst, + const HloInstruction* param) { + if (inst->opcode() == HloOpcode::kGetTupleElement && + inst->operand(0) == param) { + return inst->tuple_index(); + } + return std::nullopt; + } + + // Collects the inputs that reach some body output other than their own; a + // self update does not count. + void AccumulateOutputDependencies() { + const HloInstruction* while_input = while_body_->parameter_instruction(0); + for (int64_t i = 0; i < while_body_root_->operand_count(); ++i) { + const HloInstruction* output = while_body_root_->operand(i); + if (const std::optional index = + ParameterIndex(output, while_input)) { + // The output copies input *index; the row holds only that bit, and a + // pass through (*index == i) is a self update. + if (*index != i && candidate_bit_[*index] >= 0) { + input_deps_.Set(affecting_others_row_, candidate_bit_[*index]); + } + continue; + } + const int output_slot = SlotOf(output); + if (candidate_bit_[i] >= 0) { + input_deps_.OrIgnoringBit(affecting_others_row_, output_slot, + candidate_bit_[i]); + } else { + input_deps_.Or(affecting_others_row_, output_slot); + } + } + } + + const HloComputation* const while_body_; + const HloComputation* const while_cond_; + const HloInstruction* const while_body_root_; + const absl::Span candidate_bit_; + const int body_slots_; + const int num_slots_; + // Two extra rows of input_deps_: the inputs that reach a side effect or the + // condition, and the inputs that reach another output. + const int side_effecting_row_; + const int affecting_others_row_; + DenseBitSets input_deps_; + DenseUnionFind sets_; +}; + +// Clears index i of the surviving indices; returns 1 when it was still set. +int64_t EraseIndex(int64_t i, InlinedBitSet<>& used_tuple_indices) { + if (!used_tuple_indices.Test(i)) { + return 0; + } + used_tuple_indices.Clear(i); + return 1; +} + +// Case 1: an index that is unused after the loop can go on its own when no +// other output, side effect or condition depends on its input. Returns the +// number of indices newly removed from used_tuple_indices. +int64_t EraseIndicesDeadOnTheirOwn(const InlinedBitSet<>& used_after_loop, + const WhileInputDependencies& deps, + InlinedBitSet<>& used_tuple_indices) { + int64_t num_erased = 0; + for (int64_t i = 0; i < deps.tuple_size(); ++i) { + if (!used_after_loop.Test(i) && !deps.AffectsOtherOutputs(i) && + !deps.ReachesSink(i)) { + VLOG(2) << "Remove with dependencies " << i; + num_erased += EraseIndex(i, used_tuple_indices); + } + } + return num_erased; +} + +// Case 2: a connected group of indices can go when none of its indices blocks +// it. An index blocks its group when it is not a candidate, that is used after +// the loop and not passed through, or when a side effect or the condition +// depends on it. Returns the number of indices newly removed from +// used_tuple_indices. +int64_t EraseIndicesDeadAsGroup(WhileInputDependencies& deps, + InlinedBitSet<>& used_tuple_indices) { + InlinedBitSet<> blocked_components(deps.num_slots()); + for (int64_t i = 0; i < deps.tuple_size(); ++i) { + if (!deps.IsCandidate(i) || deps.ReachesSink(i)) { + blocked_components.Set(deps.ComponentOf(i)); + } + } + int64_t num_erased = 0; + for (int64_t i = 0; i < deps.tuple_size(); ++i) { + if (!blocked_components.Test(deps.ComponentOf(i))) { + VLOG(2) << "Remove with groups: index " << i; + num_erased += EraseIndex(i, used_tuple_indices); + } + } + return num_erased; +} + } // namespace absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { @@ -372,7 +624,6 @@ absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { } HloInstruction* while_init = while_op->mutable_operand(0); - HloComputation* while_cond = while_op->while_condition(); HloComputation* while_body = while_op->while_body(); HloInstruction* while_body_root = while_body->root_instruction(); @@ -388,12 +639,6 @@ absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { const int64_t tuple_size = ShapeUtil::TupleElementCount(while_init->shape()); auto print_no_metadata = HloPrintOptions().set_print_metadata(false); - // A set that stores all unused indices. Initialize to all indices and then - // remove elements from it. - absl::flat_hash_set used_tuple_indices; - for (int64_t i = 0; i < tuple_size; ++i) { - used_tuple_indices.insert(i); - } // Bail if param0 of while_cond or while_body has users which aren't of type // get-tuple-element. @@ -406,146 +651,49 @@ absl::StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { "empty."; return false; } - InlinedBitSet<> used_indices_after_loop(tuple_size); - if (while_op == while_op->parent()->root_instruction()) { - used_indices_after_loop.SetAll(tuple_size); - } - for (auto user : while_op->users()) { - if (user->opcode() != HloOpcode::kGetTupleElement) { - used_indices_after_loop.SetAll(tuple_size); - break; - } - used_indices_after_loop.Set(user->tuple_index()); - } + const InlinedBitSet<> used_after_loop = + TupleIndicesUsedAfterLoop(while_op, tuple_size); // We identify dead loop parameters in two cases (provided they do not cause // side effects or affect the loop condition): // 1) There is no use after the loop, and the input does not affect other // outputs (though it may update itself). // 2) A group of mutually dependent elements whose outputs are either unused - // after the loop or are passed-through unmodified. We use UnionFind to - // approximate connected components (which has false negatives for directed - // dependencies). - - // Tracks the set of inputs that each instruction depends on (in one - // iteration). Used to find inputs affecting other outputs (case 1) and - // inputs affecting the condition or side effects (cases 1 & 2). - absl::flat_hash_map> inst_input_deps; - // Find disjoint sets of connected instruction groups. This helps finding a - // group of inter-dependent indices that can be removed together. For case 2). - absl::flat_hash_map> - disjoint_sets; - // Initialize. - for (HloComputation* comp : {while_body, while_cond}) { - HloInstruction* while_input = comp->parameter_instruction(0); - for (HloInstruction* inst : comp->instructions()) { - if (inst == while_input || inst == while_body_root) { - continue; - } - disjoint_sets[inst].Get() = inst; - } - } - // Track the dependencies and merge the disjoint sets. - InlinedBitSet<> side_effecting_indices(tuple_size); - for (HloComputation* comp : {while_body, while_cond}) { - HloInstruction* while_input = comp->parameter_instruction(0); - for (HloInstruction* inst : comp->MakeInstructionPostOrder()) { - if (inst == while_input || inst == while_body_root) { - continue; - } - auto& deps = inst_input_deps.try_emplace(inst, tuple_size).first->second; - auto& my_set = disjoint_sets[inst]; - if (inst->opcode() == HloOpcode::kGetTupleElement && - inst->operand(0) == while_input) { - deps.Set(inst->tuple_index()); - HloInstruction* output = - while_body_root->mutable_operand(inst->tuple_index()); - if (output != inst) { - disjoint_sets[output].Merge(&my_set); - } - } else { - for (HloInstruction* operand : inst->operands()) { - if (operand == while_body_root) { - // The root is skipped above, so it has no map entries. Looking it - // up here would insert one and could dangle deps and my_set. - continue; - } - disjoint_sets[operand].Merge(&my_set); - auto it = inst_input_deps.find(operand); - if (it != inst_input_deps.end()) { - deps |= it->second; - } - } - } - if (inst->HasSideEffect() || inst == while_cond->root_instruction()) { - side_effecting_indices |= deps; - } - } - } - // Find inputs that can be removed because they don't affect others. - InlinedBitSet<> indices_affecting_others(tuple_size); - for (int64_t i = 0; i < tuple_size; ++i) { - HloInstruction* output = while_body_root->mutable_operand(i); - auto it = inst_input_deps.find(output); - if (it != inst_input_deps.end()) { - InlinedBitSet<> deps = it->second; - // Ignore self-updates. - deps.Clear(i); - indices_affecting_others |= deps; - } - } + // after the loop or are passed through unmodified. We use a union find data + // structure to approximate connected components (which has false negatives + // for directed dependencies). + + // Only an index that is unused after the loop or passed through unmodified + // can be removed, in either case. The dependency analysis tracks just those + // candidate indices, numbered by candidate_bit; every other index keeps -1 + // there. + std::vector candidate_bit(tuple_size, -1); + int num_candidates = 0; for (int64_t i = 0; i < tuple_size; ++i) { - if (!indices_affecting_others.Test(i) && !used_indices_after_loop.Test(i) && - !side_effecting_indices.Test(i)) { - VLOG(2) << "Remove with dependencies " << i; - used_tuple_indices.erase(i); + if (!used_after_loop.Test(i) || IsPassedThroughUnmodified(while_body, i)) { + candidate_bit[i] = num_candidates++; } } - // Find the connected groups of input/output indices. - absl::flat_hash_map> groups; - for (int64_t i = 0; i < tuple_size; ++i) { - HloInstruction* output = while_body_root->mutable_operand(i); - groups[disjoint_sets[output].Get()].insert(i); - } - for (HloComputation* comp : {while_body, while_cond}) { - HloInstruction* while_input = comp->parameter_instruction(0); - for (HloInstruction* gte : while_input->users()) { - groups[disjoint_sets[gte].Get()].insert(gte->tuple_index()); - } - } - for (const auto& group : groups) { - if (absl::c_any_of(group.second, [&](int64_t index) { - // We cannot remove this index causes side effects, or if its output - // is not passed through from input and it is used after the while op. - const HloInstruction* output = while_body_root->operand(index); - return side_effecting_indices.Test(index) || - (used_indices_after_loop.Test(index) && - !(output->opcode() == HloOpcode::kGetTupleElement && - output->operand(0) == - while_body->parameter_instruction(0) && - output->tuple_index() == index)); - })) { - continue; - } - VLOG(2) << "Remove with groups:"; - for (int64_t index : group.second) { - VLOG(2) << " index " << index; - used_tuple_indices.erase(index); - } + if (num_candidates == 0) { + VLOG(2) << "Loop " << while_op->ToString(print_no_metadata) + << " has no removable candidate index."; + return false; } + WhileInputDependencies deps(while_op, candidate_bit, num_candidates); + // The indices that survive: all of them until a case removes some. + InlinedBitSet<> used_tuple_indices(tuple_size); + used_tuple_indices.SetAll(tuple_size); + const int64_t num_removed = + EraseIndicesDeadOnTheirOwn(used_after_loop, deps, used_tuple_indices) + + EraseIndicesDeadAsGroup(deps, used_tuple_indices); - if (used_tuple_indices.size() == tuple_size) { + if (num_removed == 0) { VLOG(2) << "Loop " << while_op->ToString(print_no_metadata) << " uses all of its inputs; no simplification possible."; return false; } - // If we got here, used_tuple_indices.size() < tuple_size, meaning some - // elements of the loop's tuple aren't used by while_body or while_cond. - CHECK_LT(used_tuple_indices.size(), tuple_size); - - VLOG(1) << "Eliminating " << tuple_size - used_tuple_indices.size() - << " elements from tuple of " + VLOG(1) << "Eliminating " << num_removed << " elements from tuple of " << while_op->ToString(print_no_metadata); ABSL_ASSIGN_OR_RETURN(while_op, @@ -627,12 +775,12 @@ absl::StatusOr RemoveRepeatedWhileTupleIndices( } // And finally, remove the dead indices. - absl::flat_hash_set used_tuple_indices; - for (int index = 0; - index < while_body->root_instruction()->shape().tuple_shapes().size(); - ++index) { - if (dead_to_surviving_index.find(index) == dead_to_surviving_index.end()) { - used_tuple_indices.insert(index); + const int64_t tuple_size = + while_body->root_instruction()->shape().tuple_shapes().size(); + InlinedBitSet<> used_tuple_indices(tuple_size); + for (int64_t index = 0; index < tuple_size; ++index) { + if (!dead_to_surviving_index.contains(index)) { + used_tuple_indices.Set(index); } } diff --git a/third_party/xla/xla/service/while_loop_simplifier_test.cc b/third_party/xla/xla/service/while_loop_simplifier_test.cc index 37d7f9a1a3c853..dc348d5027c9c4 100644 --- a/third_party/xla/xla/service/while_loop_simplifier_test.cc +++ b/third_party/xla/xla/service/while_loop_simplifier_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "absl/status/status_matchers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_replace.h" +#include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" @@ -45,6 +46,9 @@ namespace { using ::testing::_; namespace op = xla::testing::opcode_matchers; +using op::Constant; +using op::Tuple; +using op::While; // Returns the first kWhile instruction within m's entry computation. HloInstruction* FindFirstWhile(HloModule* m) { @@ -2342,5 +2346,70 @@ TEST_F(WhileLoopSimplifierTest, SimplifierWithDisabledWhileLoopDceAttr) { EXPECT_FALSE(changed); } +// An async-update or async-done carries no called computation of its own; its +// side effect is the wrapped instruction's, reached through the chain. The +// loop below feeds x to the async-start and y to the async-update only, and +// SIDE_EFFECT selects whether the wrapped custom call has a side effect. +constexpr absl::string_view kAsyncChainHlo = R"hlo( + HloModule AsyncChain + + async_comp { + a = f32[] parameter(0) + b = f32[] parameter(1) + ROOT cc = f32[] custom-call(a, b), custom_call_target="Foo", custom_call_has_side_effect=SIDE_EFFECT + } + cond { + p = (s32[], f32[], f32[]) parameter(0) + i = s32[] get-tuple-element(p), index=0 + ten = s32[] constant(10) + ROOT lt = pred[] compare(i, ten), direction=LT + } + body { + p = (s32[], f32[], f32[]) parameter(0) + i = s32[] get-tuple-element(p), index=0 + x = f32[] get-tuple-element(p), index=1 + y = f32[] get-tuple-element(p), index=2 + one = s32[] constant(1) + i.next = s32[] add(i, one) + start = ((f32[]), f32[], s32[]) async-start(x), calls=async_comp + update = ((f32[], f32[]), f32[], s32[]) async-update(start, y) + done = f32[] async-done(update) + ROOT t = (s32[], f32[], f32[]) tuple(i.next, x, y) + } + ENTRY main { + zero = s32[] constant(0) + one = f32[] constant(1) + two = f32[] constant(2) + init = (s32[], f32[], f32[]) tuple(zero, one, two) + loop = (s32[], f32[], f32[]) while(init), condition=cond, body=body + ROOT r = s32[] get-tuple-element(loop), index=0 + } + )hlo"; + +// A side effect behind the async chain keeps x and y alive, as it would +// behind the start itself. +TEST_F(WhileLoopSimplifierTest, AsyncChainWithSideEffectKeepsItsInputsAlive) { + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(absl::StrReplaceAll( + kAsyncChainHlo, {{"SIDE_EFFECT", "true"}}))); + HloInstruction* loop = FindFirstWhile(m.get()); + + ASSERT_OK_AND_ASSIGN(bool changed, TryRemoveDeadWhileParams(loop)); + + EXPECT_FALSE(changed); +} + +// Without the side effect both x and y are dead. +TEST_F(WhileLoopSimplifierTest, AsyncChainWithoutSideEffectIsDead) { + ASSERT_OK_AND_ASSIGN(auto m, + ParseAndReturnVerifiedModule(absl::StrReplaceAll( + kAsyncChainHlo, {{"SIDE_EFFECT", "false"}}))); + HloInstruction* loop = FindFirstWhile(m.get()); + + ASSERT_OK_AND_ASSIGN(bool changed, TryRemoveDeadWhileParams(loop)); + + EXPECT_TRUE(changed); + EXPECT_THAT(FindFirstWhile(m.get()), While(Tuple(Constant()))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/union_find.h b/third_party/xla/xla/union_find.h index c3001d9fdcb5fb..2f241a3737af13 100644 --- a/third_party/xla/xla/union_find.h +++ b/third_party/xla/xla/union_find.h @@ -16,6 +16,12 @@ limitations under the License. #ifndef XLA_UNION_FIND_H_ #define XLA_UNION_FIND_H_ +#include +#include +#include + +#include "absl/log/check.h" + namespace xla { // Union-Find data structure. @@ -77,6 +83,75 @@ UnionFind* UnionFind::FindRoot() { return parent_; } +// Union-Find optimized for a dense index range [0, num_elements). Use it +// instead of UnionFind when the elements already have dense ids (for +// example the local ids of a computation's instructions). Avoids per element +// allocations and the need for a hash_map from element to cluster: the whole +// structure is two flat arrays. +class DenseUnionFind { + public: + // Every element starts in its own cluster: both arrays zero initialized. + explicit DenseUnionFind(int num_elements) + : parent_plus_one_(CheckNonNegative(num_elements), 0), + rank_(num_elements, 0) {} + + int num_elements() const { return parent_plus_one_.size(); } + + // Returns the representative of the cluster containing 'element'. Compresses + // the path on the way up by path halving, the one pass variant of path + // compression: every other node on the path is pointed at its grandparent. + int Find(int element) { + DCHECK_GE(element, 0); + DCHECK_LT(element, num_elements()); + while (true) { + const int parent = Parent(element); + if (parent == element) { + return element; + } + const int grandparent = Parent(parent); + parent_plus_one_[element] = grandparent + 1; + element = grandparent; + } + } + + // Merges the clusters of 'a' and 'b' by rank: the root of higher rank + // becomes the representative of the merged cluster, and equal ranks grow the + // surviving root's rank by one. + void Merge(int a, int b) { + a = Find(a); + b = Find(b); + if (a == b) { + return; + } + if (rank_[a] < rank_[b]) { + std::swap(a, b); + } + if (rank_[a] == rank_[b]) { + ++rank_[a]; + } + parent_plus_one_[b] = a + 1; + } + + private: + static int CheckNonNegative(int count) { + DCHECK_GE(count, 0); + return count; + } + + // A root is its own parent. + int Parent(int element) const { + const int stored = parent_plus_one_[element]; + return stored == 0 ? element : stored - 1; + } + + // The parent of each element plus one; 0 marks a root, so the zero + // initialized array holds singletons and needs no iota. + std::vector parent_plus_one_; + // The rank of a root bounds the height of its tree and stays below + // log2(num_elements), so a byte holds it. + std::vector rank_; +}; + } // namespace xla #endif // XLA_UNION_FIND_H_ diff --git a/third_party/xla/xla/union_find_test.cc b/third_party/xla/xla/union_find_test.cc new file mode 100644 index 00000000000000..331be60ec12dd8 --- /dev/null +++ b/third_party/xla/xla/union_find_test.cc @@ -0,0 +1,209 @@ +/* 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/union_find.h" + +#include +#include + +#include "xla/tsl/platform/test.h" +#include "xla/tsl/platform/test_benchmark.h" + +namespace xla { +namespace { + +using ::benchmark::DoNotOptimize; +using ::testing::benchmark::State; + +TEST(UnionFindTest, MergeKeepsTheCallersValue) { + UnionFind a(1); + UnionFind b(2); + UnionFind c(3); + + a.Merge(&b); + c.Merge(&a); + + EXPECT_EQ(a.Get(), 3); + EXPECT_EQ(b.Get(), 3); + EXPECT_EQ(c.Get(), 3); + EXPECT_EQ(a.Size(), 3); +} + +TEST(DenseUnionFindTest, ZeroElements) { + DenseUnionFind sets(0); + EXPECT_EQ(sets.num_elements(), 0); +} + +TEST(DenseUnionFindTest, StartsAsSingletons) { + DenseUnionFind sets(4); + + EXPECT_EQ(sets.num_elements(), 4); + for (int i = 0; i < 4; ++i) { + EXPECT_EQ(sets.Find(i), i); + } +} + +TEST(DenseUnionFindTest, MergeJoinsClusters) { + DenseUnionFind sets(6); + + sets.Merge(0, 1); + sets.Merge(2, 3); + sets.Merge(1, 3); + sets.Merge(4, 4); + + EXPECT_EQ(sets.Find(0), sets.Find(3)); + EXPECT_EQ(sets.Find(1), sets.Find(2)); + EXPECT_NE(sets.Find(0), sets.Find(4)); + EXPECT_NE(sets.Find(4), sets.Find(5)); + EXPECT_EQ(sets.Find(4), 4); + EXPECT_EQ(sets.Find(5), 5); +} + +// Union by rank: the root of higher rank represents the merged cluster, and a +// tie keeps the first argument's root and raises its rank. +TEST(DenseUnionFindTest, RepresentativeIsTheRootOfHigherRank) { + DenseUnionFind sets(5); + + sets.Merge(0, 1); // Tie: 0 stays the root, now of rank 1. + EXPECT_EQ(sets.Find(1), 0); + sets.Merge(2, 0); // 2 has rank 0, so 0 keeps representing the cluster. + EXPECT_EQ(sets.Find(2), 0); + sets.Merge(3, 4); // Tie: 3 becomes a root of rank 1. + sets.Merge(3, 0); // Tie between ranks 1: 3 stays the root, now of rank 2. + for (int i = 0; i < 5; ++i) { + EXPECT_EQ(sets.Find(i), 3); + } +} + +// Reference implementation: one label per element, relabelled on every merge. +class BruteForceSets { + public: + explicit BruteForceSets(int n) : label_(n) { + for (int i = 0; i < n; ++i) { + label_[i] = i; + } + } + bool SameSet(int a, int b) const { return label_[a] == label_[b]; } + void Merge(int a, int b) { + const int from = label_[b]; + const int to = label_[a]; + for (int& label : label_) { + if (label == from) { + label = to; + } + } + } + + private: + std::vector label_; +}; + +TEST(DenseUnionFindTest, MatchesBruteForceOnRandomMerges) { + constexpr int kNumInstances = 100; + constexpr int kNumMerges = 200; + std::mt19937 rng(11); + std::uniform_int_distribution size_dist(1, 60); + + for (int instance = 0; instance < kNumInstances; ++instance) { + const int n = size_dist(rng); + std::uniform_int_distribution element(0, n - 1); + DenseUnionFind sets(n); + BruteForceSets reference(n); + + for (int merge = 0; merge < kNumMerges; ++merge) { + const int a = element(rng); + const int b = element(rng); + sets.Merge(a, b); + reference.Merge(a, b); + const int x = element(rng); + const int y = element(rng); + ASSERT_EQ(sets.Find(x) == sets.Find(y), reference.SameSet(x, y)) + << "instance " << instance << " merge " << merge; + } + for (int x = 0; x < n; ++x) { + for (int y = 0; y < n; ++y) { + ASSERT_EQ(sets.Find(x) == sets.Find(y), reference.SameSet(x, y)) + << "instance " << instance << " pair " << x << ", " << y; + } + } + } +} + +// The access pattern of TryRemoveDeadWhileParams over the instruction slots of +// a loop, including the construction of the structure: +// - one merge per operand edge, +// - then one find per tuple index. +std::vector RandomEdges(int n) { + std::mt19937 rng(5); + std::uniform_int_distribution element(0, n - 1); + std::vector edges(2 * n); + for (int& e : edges) { + e = element(rng); + } + return edges; +} + +void BM_DenseUnionFindTryRemoveDeadWhileParamsPattern(State& state) { + const int n = state.range(0); + const std::vector edges = RandomEdges(n); + for (auto s : state) { + DenseUnionFind sets(n); + for (int i = 0; i + 1 < edges.size(); i += 2) { + sets.Merge(edges[i], edges[i + 1]); + } + int sink = 0; + for (int i = 0; i < n; ++i) { + sink += sets.Find(i); + } + DoNotOptimize(sink); + } + state.SetItemsProcessed(state.iterations() * 2 * n); +} +BENCHMARK(BM_DenseUnionFindTryRemoveDeadWhileParamsPattern) + ->Arg(1024) + ->Arg(100000) + ->Arg(1 << 20) + ->Arg(1 << 22); + +// The same pattern on UnionFind: one node per element, holding its index as +// the value, so a find is a Get. +void BM_UnionFindTryRemoveDeadWhileParamsPattern(State& state) { + const int n = state.range(0); + const std::vector edges = RandomEdges(n); + for (auto s : state) { + std::vector> sets; + sets.reserve(n); + for (int i = 0; i < n; ++i) { + sets.emplace_back(i); + } + for (int i = 0; i + 1 < edges.size(); i += 2) { + sets[edges[i]].Merge(&sets[edges[i + 1]]); + } + int sink = 0; + for (int i = 0; i < n; ++i) { + sink += sets[i].Get(); + } + DoNotOptimize(sink); + } + state.SetItemsProcessed(state.iterations() * 2 * n); +} +BENCHMARK(BM_UnionFindTryRemoveDeadWhileParamsPattern) + ->Arg(1024) + ->Arg(100000) + ->Arg(1 << 20) + ->Arg(1 << 22); + +} // namespace +} // namespace xla From 0a82ab256d7d5c5c5fc18c2f3c4b5654da5225bf Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 29 Sep 2026 09:56:33 -0700 Subject: [PATCH 22/25] Move JAX collective microbenchmarks to xla/benchmarks/collectives and integrate with the benchmark driver. Move collectives_benchmark.py and all_gather_benchmark.py from internal benchmarks to xla/benchmarks/collectives/, add CollectiveBenchmarkConfig and CollectiveBenchmark adapters, and register 'collectives' in BENCHMARK_FACTORIES so that run_benchmarks.py --benchmarks=collectives runs the collective microbenchmark suite and outputs CSV metrics. PiperOrigin-RevId: 990370997 --- third_party/xla/xla/benchmarks/README.md | 3 + .../xla/xla/benchmarks/benchmark_configs.py | 25 + .../xla/benchmarks/benchmark_configs_test.py | 73 +++ .../collectives/all_gather_benchmark.py | 160 +++++++ .../collectives/collectives_benchmark.py | 452 ++++++++++++++++++ 5 files changed, 713 insertions(+) create mode 100644 third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py create mode 100644 third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py diff --git a/third_party/xla/xla/benchmarks/README.md b/third_party/xla/xla/benchmarks/README.md index 484da58bc44987..e3efcc33d4cdbc 100644 --- a/third_party/xla/xla/benchmarks/README.md +++ b/third_party/xla/xla/benchmarks/README.md @@ -24,6 +24,9 @@ python3 xla/benchmarks/run_benchmarks.py --benchmarks=dense_matmul --csv_path=

+# Run collectives benchmark suite and write results to a CSV file +python3 xla/benchmarks/run_benchmarks.py --benchmarks=collectives --csv_path= + # Run full benchmark suite and write results to multiple CSV files python3 xla/benchmarks/run_benchmarks.py --csv_path= ``` diff --git a/third_party/xla/xla/benchmarks/benchmark_configs.py b/third_party/xla/xla/benchmarks/benchmark_configs.py index d13b0155b4acdf..fe46e6393cc326 100644 --- a/third_party/xla/xla/benchmarks/benchmark_configs.py +++ b/third_party/xla/xla/benchmarks/benchmark_configs.py @@ -21,6 +21,8 @@ from jax.experimental.pallas import tpu as pltpu import jax.numpy as jnp +from xla.benchmarks.collectives import all_gather_benchmark +from xla.benchmarks.collectives import collectives_benchmark from xla.benchmarks.core import platform_info from xla.benchmarks.dma_microbenchmarks import chip_to_chip_dma_benchmark from xla.benchmarks.dma_microbenchmarks import chiplet_to_chiplet_dma_benchmark @@ -232,6 +234,28 @@ def get_dma_configs( return configs +_COLLECTIVE_SUITES: tuple[ + tuple[str, type[collectives_benchmark.CollectivesBenchmarks]], ... +] = (("all_gather", all_gather_benchmark.AllGatherBenchmarks),) + + +def get_collectives_configs( + chip_version: pltpu.ChipVersion | None = None, +) -> list[collectives_benchmark.CollectiveBenchmarkConfig]: + """Generates configs for all collective microbenchmarks.""" + del chip_version # Unused. + configs = [] + for suite, test_class in _COLLECTIVE_SUITES: + for name in sorted(dir(test_class)): + if name.startswith("test_"): + configs.append( + collectives_benchmark.CollectiveBenchmarkConfig( + suite=suite, test_class=test_class, test_name=name + ) + ) + return configs + + BENCHMARK_FACTORIES: Mapping[ str, Callable[[pltpu.ChipVersion | None], list[Any]] ] = immutabledict.immutabledict({ @@ -239,4 +263,5 @@ def get_dma_configs( "subchannel_matmul": get_subchannel_matmul_configs, "jax_matmul": get_jax_matmul_configs, "dma": get_dma_configs, + "collectives": get_collectives_configs, }) diff --git a/third_party/xla/xla/benchmarks/benchmark_configs_test.py b/third_party/xla/xla/benchmarks/benchmark_configs_test.py index acfa18f988de10..638ee240c8205f 100644 --- a/third_party/xla/xla/benchmarks/benchmark_configs_test.py +++ b/third_party/xla/xla/benchmarks/benchmark_configs_test.py @@ -23,6 +23,7 @@ from xla.benchmarks import benchmark_configs from xla.benchmarks import results_utils +from xla.benchmarks.collectives import collectives_benchmark from xla.benchmarks.dma_microbenchmarks import memory_base from xla.benchmarks.jax_microbenchmarks import matmul_lib from xla.benchmarks.pallas_microbenchmarks import dense_matmul_lib @@ -214,6 +215,78 @@ def test_bw(self): self.assertAlmostEqual(df["latency_us"].iloc[0], 15.0) self.assertGreater(df["bandwidth_gbps"].iloc[0], 0.0) + def test_collectives_configs(self): + configs = benchmark_configs.get_collectives_configs() + self.assertLen(configs, 2) + for cfg in configs: + self.assertIsInstance( + cfg, collectives_benchmark.CollectiveBenchmarkConfig + ) + self.assertEqual(cfg.suite, "all_gather") + self.assertStartsWith(cfg.test_name, "test_all_gather_") + self.assertEqual( + cfg.as_dict(), {"suite": cfg.suite, "test": cfg.test_name} + ) + + def test_collectives_benchmark_end_to_end_and_skip(self): + # 1. On CPU without TPU, CollectiveBenchmark.run() catches SkipTest. + cfg_skip = benchmark_configs.get_collectives_configs()[0] + self.assertEqual(cfg_skip.get_benchmark().run(), [None]) + + # 2. With TPU info patched using spec_set=True, run() executes the test + # method, records bandwidth/latency metrics, and populates the table. + class _FakeCollectiveSuite(collectives_benchmark.CollectivesBenchmarks): + + def setUp(self): + absltest.TestCase.setUp(self) + self.latencies_us = [] + self.metrics = {} + + def test_ag(self): + self._print_collective_bandwidth_statistics( + op_name="all-gather-major-dim", + per_device_size_bytes=256 * 1024 * 1024, + total_size_bytes=1024 * 1024 * 1024, + bus_data_bytes=768 * 1024 * 1024, + num_devices=4, + latencies_us=[100.0, 200.0], + ) + + cfg = collectives_benchmark.CollectiveBenchmarkConfig( + suite="all_gather", test_class=_FakeCollectiveSuite, test_name="test_ag" + ) + with unittest.mock.patch.object( + pltpu, + "get_tpu_info", + spec_set=True, + return_value=pltpu.get_tpu_info_for_chip(pltpu.ChipVersion.TPU_V5E, 1), + ): + prof_results = cfg.get_benchmark().run() + + df = results_utils.create_results_table([(cfg, prof_results)]) + self.assertEqual( + list(df.columns), + [ + "suite", + "test", + "per_device_size_mib", + "total_size_mib", + "num_devices", + "bus_bandwidth_gbps", + "effective_bandwidth_gbps", + "latency_us", + "flops", + ], + ) + self.assertEqual(df["suite"].iloc[0], "all_gather") + self.assertEqual(df["test"].iloc[0], "test_ag") + self.assertEqual(df["per_device_size_mib"].iloc[0], 256.0) + self.assertEqual(df["total_size_mib"].iloc[0], 1024.0) + self.assertEqual(df["num_devices"].iloc[0], 4) + self.assertAlmostEqual(df["latency_us"].iloc[0], 150.0) + self.assertGreater(df["bus_bandwidth_gbps"].iloc[0], 0.0) + self.assertGreater(df["effective_bandwidth_gbps"].iloc[0], 0.0) + if __name__ == "__main__": absltest.main() diff --git a/third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py b/third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py new file mode 100644 index 00000000000000..5eccc95260ae88 --- /dev/null +++ b/third_party/xla/xla/benchmarks/collectives/all_gather_benchmark.py @@ -0,0 +1,160 @@ +# 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. +"""JAX microbenchmark for TPU all-gather collective bandwidth and latency.""" + +from absl.testing import absltest +import jax +from jax import numpy as jnp +from jax.experimental import shard_map +import numpy as np + +from xla.benchmarks.collectives import collectives_benchmark + + +def zero_crop(x): + """Prevents x from being a live-out buffer by passing through a ZeroCrop FFI call.""" + if jax.default_backend() == "cpu": + return x + return jax.ffi.ffi_call( + "ZeroCrop", + result_shape_dtypes=jax.ShapeDtypeStruct(x.shape, x.dtype), + has_side_effect=True, + )(x) + + +class AllGatherBenchmarks(collectives_benchmark.CollectivesBenchmarks): + """JAX benchmarks for measuring single-tray all-gather bandwidth and latency.""" + + def _setup_mesh_and_axis(self): + """Creates the single-tray device mesh and returns (mesh, ag_axis, num_devices).""" + devices = self.get_single_tray_devices() + num_devices = len(devices) + + num_chips = len({tuple(d.coords) for d in devices}) + cores_per_chip = num_devices // num_chips + + if cores_per_chip > 1: + mesh = jax.sharding.Mesh( + np.array(devices).reshape((num_chips, cores_per_chip)), + ("chips", "cores"), + ) + ag_axis = ("chips", "cores") + else: + mesh = jax.sharding.Mesh(np.array(devices), ("chips",)) + ag_axis = "chips" + + return mesh, ag_axis, num_devices + + def _run_all_gather_benchmark( + self, + global_shape, + in_partition, + gather_axis, + op_name, + ): + """Runs a single-HLO all-gather benchmark and logs bandwidth/latency stats.""" + mesh, ag_axis, num_devices = self._setup_mesh_and_axis() + out_partition = jax.sharding.PartitionSpec(*([None] * len(global_shape))) + + sharding = jax.sharding.NamedSharding(mesh, in_partition) + arr = jax.random.uniform( + jax.random.key(0), + global_shape, + dtype=jnp.bfloat16, + ) + array_on_devices = jax.device_put(arr, sharding) + per_device_bytes = arr.nbytes // num_devices + + @jax.jit + def all_gather_fn(x): + return shard_map.shard_map( + lambda x_shard: zero_crop( + jax.lax.all_gather(x_shard, ag_axis, axis=gather_axis) + ), + mesh=mesh, + in_specs=in_partition, + out_specs=out_partition, + check_rep=False, + )(x) + + # Warmup run to compile and initialize buffers before tracing. + warmup_res = all_gather_fn(array_on_devices) + jax.block_until_ready(warmup_res) + + execution_times_us = self._measure_execution_times( + all_gather_fn, + ["all-gather"], + array_on_devices, + ) + self._print_all_gather_bandwidth_statistics( + per_device_size_bytes=per_device_bytes, + num_devices=num_devices, + latencies_us=execution_times_us, + op_name=op_name, + ) + + def _print_all_gather_bandwidth_statistics( + self, + per_device_size_bytes, + num_devices, + latencies_us, + op_name="all-gather", + ): + """Computes all-gather data sizes and prints bandwidth/latency statistics.""" + total_gathered_bytes = per_device_size_bytes * num_devices + # In ring-based all-gather, each device transmits (N-1)/N of the + # gathered data. + bus_data_bytes = per_device_size_bytes * (num_devices - 1) + self._print_collective_bandwidth_statistics( + op_name=op_name, + per_device_size_bytes=per_device_size_bytes, + total_size_bytes=total_gathered_bytes, + bus_data_bytes=bus_data_bytes, + num_devices=num_devices, + latencies_us=latencies_us, + ) + + def _get_per_device_chunk_mib(self, num_devices): + """Returns per-device chunk size in MiB from collective_size_mib and num_devices.""" + if self.collective_size_mib % num_devices != 0: + raise ValueError( + f"collective_size_mib ({self.collective_size_mib}) must be divisible" + f" by num_devices ({num_devices})." + ) + return self.collective_size_mib // num_devices + + def test_all_gather_major_dim_bandwidth(self): + _, ag_axis, num_devices = self._setup_mesh_and_axis() + chunk_size_mib = self._get_per_device_chunk_mib(num_devices) + self._run_all_gather_benchmark( + global_shape=(num_devices, chunk_size_mib * 512, 1024), + in_partition=jax.sharding.PartitionSpec(ag_axis, None, None), + gather_axis=0, + op_name="all-gather-major-dim", + ) + + def test_all_gather_minor_dim_bandwidth(self): + _, ag_axis, num_devices = self._setup_mesh_and_axis() + chunk_size_mib = self._get_per_device_chunk_mib(num_devices) + self._run_all_gather_benchmark( + global_shape=(chunk_size_mib, num_devices * 512, 1024), + in_partition=jax.sharding.PartitionSpec(None, ag_axis, None), + gather_axis=1, + op_name="all-gather-minor-dim", + ) + + +if __name__ == "__main__": + jax.config.config_with_absl() + absltest.main() diff --git a/third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py b/third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py new file mode 100644 index 00000000000000..23b216829a7d86 --- /dev/null +++ b/third_party/xla/xla/benchmarks/collectives/collectives_benchmark.py @@ -0,0 +1,452 @@ +# 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. +"""Base class and utilities for JAX collective microbenchmarks.""" + +import gzip +import json +import os +import pathlib +import shutil +import sys +import tempfile + +from absl import flags +from absl import logging +from absl.testing import absltest +import immutabledict +import jax +import jax.experimental.pallas.tpu as pltpu +import numpy as np + +from xla.benchmarks.core import benchmark +from xla.benchmarks.jax_microbenchmarks import jax_profiler_utils + +_NUMBER_OF_MEASUREMENTS = flags.DEFINE_integer( + "number_of_measurements", + default=5, + help="Number of measurements to take. Default: 5", + allow_override=True, +) + +_COLLECTIVE_SIZE_MIB = flags.DEFINE_integer( + "collective_size_mib", + default=1024, + help="Total collective size in MiB across all devices. Default: 1024", +) + +EXPECTED_CHIPS_PER_HOST = immutabledict.immutabledict({ + pltpu.ChipVersion.TPU_V4: 4, + pltpu.ChipVersion.TPU_V4I: 4, + pltpu.ChipVersion.TPU_V5E: 4, + pltpu.ChipVersion.TPU_V5P: 4, + pltpu.ChipVersion.TPU_V6E: 4, + pltpu.ChipVersion.TPU_7: 4, + pltpu.ChipVersion.TPU_7X: 4, + pltpu.ChipVersion.TPU_8I: 2, + pltpu.ChipVersion.TPU_8T: 2, +}) + + +class CollectivesBenchmarks(absltest.TestCase): + """Base benchmark class for measuring JAX collective operations.""" + + def setUp(self): + super().setUp() + if not any(device.platform == "tpu" for device in jax.devices()): + self.skipTest("This test requires TPU hardware.") + self.number_of_measurements = _NUMBER_OF_MEASUREMENTS.value + self.collective_size_mib = _COLLECTIVE_SIZE_MIB.value + self.latencies_us = [] + self.metrics = {} + + def verify_single_tray_setup(self): + """Verifies that the test is running on exactly one full TPU tray.""" + if not any(device.platform == "tpu" for device in jax.devices()): + self.skipTest("This test requires TPU hardware.") + + if jax.process_count() > 1: + self.skipTest( + "Single-host test expected process_count=1, got" + f" {jax.process_count()}." + ) + + local_devices = jax.local_devices() + all_devices = jax.devices() + if len(local_devices) != len(all_devices): + self.skipTest( + f"Not all devices are local to this host: local={len(local_devices)}," + f" all={len(all_devices)}." + ) + + tpu_info = pltpu.get_tpu_info() + chip_version = tpu_info.chip_version + if not isinstance(chip_version, pltpu.ChipVersion): + logging.warning( + "Unknown TPU generation %s for chip count verification.", chip_version + ) + return + expected_chips = EXPECTED_CHIPS_PER_HOST.get(chip_version) + if expected_chips is None: + logging.warning( + "Unknown TPU generation %s for chip count verification.", chip_version + ) + return + + # Check unique chip coordinates across local TPU devices. + coords = [tuple(d.coords) for d in local_devices] + num_unique_chips = len(set(coords)) + if num_unique_chips != expected_chips: + self.skipTest( + f"Expected {expected_chips} TPU chips for a single {chip_version}" + f" tray, but found {num_unique_chips}." + ) + + # Check device count based on cores per chip. + num_cores_per_chip = chip_version.num_physical_tensor_cores_per_chip + expected_devices_split = expected_chips * num_cores_per_chip + expected_devices_megacore = expected_chips + + if len(local_devices) not in ( + expected_devices_split, + expected_devices_megacore, + ): + self.skipTest( + f"Expected {expected_devices_split} (or {expected_devices_megacore}) " + f"devices for a single {chip_version} tray, got " + f"{len(local_devices)}." + ) + + def get_single_tray_devices(self): + """Verifies the single-tray setup and returns the list of local devices.""" + self.verify_single_tray_setup() + return jax.local_devices() + + def _normalize_kernel_names(self, kernel_names): + """Normalizes kernel names into a list of strings.""" + if isinstance(kernel_names, str): + return [kernel_names] + elif isinstance(kernel_names, (tuple, list)): + return list(kernel_names) + raise TypeError( + "`kernel_names` must be a string, tuple, or list of strings. " + f"Got: {type(kernel_names)}" + ) + + def _extract_execution_times(self, profiler_dir, kernel_names): + """Extracts execution times from profiler data for matching kernel names.""" + undeclared_outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR") + if undeclared_outputs_dir: + test_name = self._testMethodName + for profile_file in pathlib.Path(profiler_dir).rglob("*"): + if profile_file.is_file() and ( + profile_file.name.endswith(".xplane.pb") + or profile_file.name.endswith(".trace.json.gz") + ): + dest = ( + pathlib.Path(undeclared_outputs_dir) + / f"{test_name}_{profile_file.name}" + ) + shutil.copy2(profile_file, dest) + logging.info( + "Saved XProf profile artifact to undeclared outputs: %s", dest + ) + + trace_files = list(pathlib.Path(profiler_dir).glob("**/*.trace.json.gz")) + if not trace_files: + raise FileNotFoundError(f"Could not find trace.json.gz in {profiler_dir}") + + kernel_durs_per_core = {} + sc_kernel_durs_per_core = {} + for trace_file in trace_files: + with gzip.open(trace_file, "rt") as f: + trace_events = json.load(f).get("traceEvents", []) + + has_sc_offload = any( + e.get("args", {}).get("hlo_category") == "async-done" + and e.get("args", {}).get("offload_type") == "OFFLOAD_COLLECTIVE" + for e in trace_events + ) + matched_count = 0 + sample_events = [] + for event in trace_events: + event_name = event.get("name", "") + args = event.get("args", {}) + tf_op = args.get("tf_op", "") + hlo_op = args.get("hlo_op", "") + if "dur" in event: + is_host_wrapper = event_name.startswith(( + "PjitFunction", + "CommonPjRt", + "PJRT_", + "PjRt", + "$", + "Mutex", + "PythonRefManager", + "ParseArguments", + "Wait for ", + )) + if len(sample_events) < 30 and event_name and not is_host_wrapper: + sample_events.append( + f"pid={event.get('pid')} name='{event_name}'" + f" tf_op='{tf_op}' hlo_op='{hlo_op}'" + ) + if is_host_wrapper: + continue + hlo_category = args.get("hlo_category", "") + if hlo_category in ("async-start", "async-update"): + continue + offload_type = args.get("offload_type", "") + is_sc_offload = ( + hlo_category == "async-done" + and offload_type == "OFFLOAD_COLLECTIVE" + ) + for kernel_name in kernel_names: + k_hyphen = kernel_name.replace("_", "-") + k_under = kernel_name.replace("-", "_") + is_name_match = event_name.startswith(( + k_hyphen, + k_under, + "ragged-all-to-all", + )) + if has_sc_offload: + if is_sc_offload: + pid = event.get("pid") + if pid not in kernel_durs_per_core: + kernel_durs_per_core[pid] = [] + kernel_durs_per_core[pid].append(event["dur"]) + matched_count += 1 + logging.info( + "[Trace Event] pid=%s, dur=%.2f us, name='%s', tf_op='%s'," + " hlo_op='%s', hlo_category='%s', offload_type='%s'", + pid, + event["dur"], + event_name, + tf_op, + hlo_op, + hlo_category, + offload_type, + ) + break + elif is_name_match and not offload_type: + pid = event.get("pid") + if pid not in sc_kernel_durs_per_core: + sc_kernel_durs_per_core[pid] = [] + sc_kernel_durs_per_core[pid].append(event["dur"]) + break + elif is_name_match: + pid = event.get("pid") + if pid not in kernel_durs_per_core: + kernel_durs_per_core[pid] = [] + kernel_durs_per_core[pid].append(event["dur"]) + matched_count += 1 + logging.info( + "[Trace Event] pid=%s, dur=%.2f us, name='%s', tf_op='%s'," + " hlo_op='%s', hlo_category='%s', offload_type='%s'", + pid, + event["dur"], + event_name, + tf_op, + hlo_op, + hlo_category, + offload_type, + ) + break + if matched_count == 0: + logging.warning( + "No events matched %s in %s. Non-host sample trace events with" + " dur: %s", + kernel_names, + trace_file.name, + sample_events, + ) + logging.info( + "Processed trace file %s: found %d matching events across %d" + " pids/cores.", + trace_file.name, + matched_count, + len(kernel_durs_per_core), + ) + + if sc_kernel_durs_per_core: + sc_durs = [] + for pid, durs in sc_kernel_durs_per_core.items(): + sc_durs.extend(durs) + if sc_durs: + logging.info( + "SparseCore pure kernel execution durations across %d PIDs (%d" + " samples): %.2f +/- %.2f us", + len(sc_kernel_durs_per_core), + len(sc_durs), + np.mean(sc_durs), + np.std(sc_durs), + ) + + execution_times_us = [] + for pid, durs in kernel_durs_per_core.items(): + if durs: + logging.info("PID %s event durations: %s", pid, durs) + num_measurements = self.number_of_measurements + if len(durs) >= num_measurements and len(durs) % num_measurements == 0: + k = len(durs) // num_measurements + for i in range(0, len(durs), k): + execution_times_us.append(sum(durs[i : i + k])) + else: + execution_times_us.extend(durs) + + if not execution_times_us: + raise ValueError( + f"Could not find execution times matching {kernel_names}." + ) + return execution_times_us + + def _measure_execution_times( + self, + kernel, + kernel_names, + *kernel_args, + **kernel_kwargs, + ): + """Collect XProf measurements using the JAX Profiler API.""" + normalized_names = self._normalize_kernel_names(kernel_names) + + # Warmup run to compile and initialize buffers before tracing. + warmup_res = kernel(*kernel_args, **kernel_kwargs) + jax.block_until_ready(warmup_res) + + with tempfile.TemporaryDirectory() as tmpdir: + with jax.profiler.trace(tmpdir): + for _ in range(self.number_of_measurements): + result = kernel(*kernel_args, **kernel_kwargs) + jax.block_until_ready(result) + return self._extract_execution_times(tmpdir, normalized_names) + + def _print_collective_bandwidth_statistics( + self, + op_name, + per_device_size_bytes, + total_size_bytes, + bus_data_bytes, + num_devices, + latencies_us, + ): + """Prints bandwidth and latency statistics for a collective operation.""" + latencies_ns = np.array(latencies_us) * 1e3 + bus_bandwidth_gbps = bus_data_bytes / latencies_ns + effective_bandwidth_gbps = total_size_bytes / latencies_ns + + avg_bus_bw = float(np.mean(bus_bandwidth_gbps)) + std_bus_bw = float(np.std(bus_bandwidth_gbps)) + avg_eff_bw = float(np.mean(effective_bandwidth_gbps)) + std_eff_bw = float(np.std(effective_bandwidth_gbps)) + avg_latency_us = float(np.mean(latencies_us)) + std_latency_us = float(np.std(latencies_us)) + + tpu_info = pltpu.get_tpu_info() + + coords = [tuple(getattr(d, "coords", (0,))) for d in jax.local_devices()] + mesh_dims = "x".join(str(max(dim) + 1) for dim in zip(*coords)) + num_chips = len(set(coords)) + cores_per_chip = num_devices // num_chips if num_chips else 1 + topology_str = ( + f"{mesh_dims} mesh ({num_chips} chips x {cores_per_chip} cores)" + ) + + flags_list = [arg for arg in sys.argv[1:] if arg.startswith("--xla_")] + flags_str = " ".join(flags_list) if flags_list else "default" + + logging.info("==================================================") + logging.info("Test: %s", self._testMethodName) + logging.info("\tCollective Op: %s", op_name) + logging.info("\tTPU generation: %s", tpu_info.chip_version) + logging.info("\tNumber of devices: %d", num_devices) + logging.info("\tTopology: %s", topology_str) + logging.info("\tFlags: %s", flags_str) + if per_device_size_bytes >= 1024**2: + logging.info( + "\tPer-device chunk size: %.2f MiB", + per_device_size_bytes / (1024**2), + ) + else: + logging.info( + "\tPer-device chunk size: %.2f KiB", per_device_size_bytes / 1024 + ) + if total_size_bytes >= 1024**2: + logging.info( + "\tTotal collective size: %.2f MiB", total_size_bytes / (1024**2) + ) + else: + logging.info( + "\tTotal collective size: %.2f KiB", total_size_bytes / 1024 + ) + logging.info("\tLatency: %.2f +/- %.2f us", avg_latency_us, std_latency_us) + logging.info("\tBus Bandwidth: %.2f +/- %.2f GB/s", avg_bus_bw, std_bus_bw) + logging.info( + "\tEffective Bandwidth: %.2f +/- %.2f GB/s", avg_eff_bw, std_eff_bw + ) + logging.info("==================================================") + + self.latencies_us = [float(x) for x in latencies_us] + self.metrics = { + "per_device_size_mib": float(per_device_size_bytes / (1024**2)), + "total_size_mib": float(total_size_bytes / (1024**2)), + "num_devices": num_devices, + "bus_bandwidth_gbps": avg_bus_bw, + "effective_bandwidth_gbps": avg_eff_bw, + } + + +class CollectiveBenchmarkConfig(benchmark.BenchmarkConfig): + """Config identifying a single collective test method for the benchmark driver.""" + + def __init__(self, suite, test_class, test_name): + self.suite = suite + self.test_class = test_class + self.test_name = test_name + self.metrics = {} + + def as_dict(self): + return {"suite": self.suite, "test": self.test_name, **self.metrics} + + def get_benchmark(self): + return CollectiveBenchmark(self) + + +class CollectiveBenchmark(benchmark.Benchmark): + """Runs a collective test method and reports its latencies to the driver.""" + + def __init__(self, config): + self._config = config + + # Required by benchmark.Benchmark ABC; unused since run() delegates to test. + def get_input_shapes_and_dtypes(self): + return [] + + def target_fn(self): + return lambda: None + + def kernel_name(self): + return self._config.test_name + + def run(self, **kwargs): + del kwargs + test = self._config.test_class(self._config.test_name) + try: + test.setUp() + getattr(test, self._config.test_name)() + except absltest.SkipTest as e: + logging.warning("Skipping %s: %s", self._config.test_name, e) + return [None] + self._config.metrics.update(test.metrics) + return [jax_profiler_utils.JaxProfilerResult(runtimes_us=test.latencies_us)] + From 17543f99daee6bb551911368b6f8f687f66fc5db Mon Sep 17 00:00:00 2001 From: Dillon Sharlet Date: Tue, 29 Sep 2026 09:59:21 -0700 Subject: [PATCH 23/25] [XLA:CPU] Add support for copy operations and degenerate layout dimensions in YNNPACK Adds support for kCopy operations in YNNPACK fusion and relaxes layout checks to tolerate out-of-order dimensions when their extents are 1. This addresses fusion breakages and performance regressions identified in https://github.com/openxla/xla/issues/49131: - When reduction pipelines use keepdims=True (e.g. pairwise reductions such as gravitational N-body or RBF kernels), layout assignment can insert reshapes and copies with layouts like f32[M, N, 1]{1,0,2}. - Previously, IsLayoutSupportedByYnn rejected any layout whose minor_to_major was not strictly descending, even if the out-of-order dimensions had extent 1 (which does not affect physical element ordering). - Furthermore, kCopy was unsupported in YNNPACK, causing LibraryRewriter to truncate fusions at copy boundaries and split the pipeline into multiple custom fusions, forcing large O(M*N) intermediate buffers to be materialized in memory. Key changes: 1. IsLayoutSupportedByYnn: Ignores extent-1 dimensions when verifying descending layouts. 2. IsCopyOpSupportedByYnn: Validates copy operations whose non-degenerate dimensions have matching layout and extents. 3. ynn_matcher.h: Registers HloOpcode::kCopy in SupportedOps() and dispatches in IsOpSupported(). 4. ynn_emitter.cc: Implements DefineCopyOp using ynn_define_copy and emits copy operations in YNN subgraphs. 5. ynn_fusion_test.cc: Adds unit tests for copy operations and degenerate layout copies across bf16, f32, and f64. PiperOrigin-RevId: 990372521 --- .../xla/backends/cpu/tests/ynn_fusion_test.cc | 37 +++++++++++++++++++ .../xla/backends/cpu/transforms/ynn_matcher.h | 5 ++- .../cpu/transforms/ynn_matcher_test.cc | 27 ++++++++++++++ .../xla/xla/backends/cpu/ynn_emitter.cc | 27 +++++++++++++- .../xla/xla/backends/cpu/ynn_support.cc | 35 +++++++++++++++++- .../xla/xla/backends/cpu/ynn_support.h | 3 ++ 6 files changed, 131 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc b/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc index 6ed7c4a6260beb..8bf18a9fb03912 100644 --- a/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc +++ b/third_party/xla/xla/backends/cpu/tests/ynn_fusion_test.cc @@ -113,6 +113,43 @@ TEST_P(YnnFusionTest, Pad) { RunTest(kModuleStr); } +TEST_P(YnnFusionTest, Copy) { + constexpr absl::string_view kModuleStr = R"( + HloModule copy + + ynn_fusion { + %input = $dtype[8, 10] parameter(0) + ROOT %copy = $dtype[8, 10] copy(%input) + } + + ENTRY entry { + %p0 = $dtype[8, 10] parameter(0) + ROOT %fusion = $dtype[8, 10] fusion(%p0), kind=kCustom, calls=ynn_fusion, + backend_config={"fusion_config": {kind: "__ynn_fusion"}} + })"; + + RunTest(kModuleStr); +} + +TEST_P(YnnFusionTest, CopyDegenerateLayout) { + constexpr absl::string_view kModuleStr = R"( + HloModule copy_degenerate_layout + + ynn_fusion { + %input = $dtype[8, 10, 1]{1,0,2} parameter(0) + ROOT %copy = $dtype[8, 10, 1]{2,1,0} copy(%input) + } + + ENTRY entry { + %p0 = $dtype[8, 10, 1]{1,0,2} parameter(0) + ROOT %fusion = $dtype[8, 10, 1]{2,1,0} fusion(%p0), kind=kCustom, + calls=ynn_fusion, + backend_config={"fusion_config": {kind: "__ynn_fusion"}} + })"; + + RunTest(kModuleStr); +} + TEST_P(YnnFusionTest, Transpose) { constexpr absl::string_view kModuleStr = R"( HloModule transpose diff --git a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h index e5d18a75601cc1..f43b25f2e60df9 100644 --- a/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h +++ b/third_party/xla/xla/backends/cpu/transforms/ynn_matcher.h @@ -50,7 +50,7 @@ class YnnMatcher : public LibraryMatcher { HloOpcode::kConvolution, HloOpcode::kReshape, HloOpcode::kBitcast, HloOpcode::kBroadcast, HloOpcode::kTranspose, HloOpcode::kPad, - HloOpcode::kIota}; + HloOpcode::kCopy, HloOpcode::kIota}; for (const auto& [op, _] : GetYnnUnaryOpMap()) { supported_ops.insert(op); } @@ -85,6 +85,9 @@ class YnnMatcher : public LibraryMatcher { if (instr->opcode() == HloOpcode::kPad) { return IsPadOpSupportedByYnn(instr); } + if (instr->opcode() == HloOpcode::kCopy) { + return IsCopyOpSupportedByYnn(instr); + } if (!IsInstructionPreferredByYnn(instr)) { // TODO: It might make sense sometimes that even though an instruction is // not preferred by YNNPACK, that we should still fuse it, if it lies 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 79afc3055657ea..6d3607737c29de 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 @@ -201,6 +201,33 @@ TEST_F(YnnReduceTest, ConvertReduce) { )"); } +TEST_F(YnnReduceTest, CopyDegenerateLayout) { + const char* hlo_text = R"( + HloModule copy_degenerate_layout + + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + + ENTRY main { + input = f32[128,64,1]{1,0,2} parameter(0) + copied = f32[128,64,1]{2,1,0} copy(input) + init = f32[] constant(0) + ROOT result = f32[128] reduce(copied, init), dimensions={1,2}, to_apply=add + } + )"; + + MatchOptimizedHlo(hlo_text, R"( + CHECK: copy + CHECK: reduce + CHECK: ENTRY + CHECK: kind=kCustom + CHECK: "kind":"__ynn_fusion" + )"); +} + struct DotTestConfig { absl::string_view lhs_dtype; absl::string_view rhs_dtype; diff --git a/third_party/xla/xla/backends/cpu/ynn_emitter.cc b/third_party/xla/xla/backends/cpu/ynn_emitter.cc index 395e8ff24177d6..25ba2f528528cc 100644 --- a/third_party/xla/xla/backends/cpu/ynn_emitter.cc +++ b/third_party/xla/xla/backends/cpu/ynn_emitter.cc @@ -217,6 +217,21 @@ absl::StatusOr DefineReshapeOp(ynn_subgraph_t subgraph, return DefineBitcastOp(subgraph, tensor_ids, instr); } +absl::StatusOr DefineCopyOp(ynn_subgraph_t subgraph, + TensorIdMap& tensor_ids, + const HloInstruction* instr) { + VLOG(3) << absl::StreamFormat("Define tensor value for copy op: %s", + instr->ToString()); + CHECK_EQ(instr->opcode(), HloOpcode::kCopy); + const HloInstruction* input = instr->operand(0); + CHECK_EQ(input->shape().element_type(), instr->shape().element_type()); + ABSL_ASSIGN_OR_RETURN(auto in, FindTensorValue(tensor_ids, input)); + ABSL_ASSIGN_OR_RETURN(auto out, DefineTensorValue(subgraph, instr)); + + YNN_RETURN_IF_ERROR(ynn_define_copy(subgraph, in, &out, /*flags=*/0)); + return out; +} + absl::StatusOr DefineTransposeOp(ynn_subgraph_t subgraph, TensorIdMap& tensor_ids, const HloInstruction* instr) { @@ -845,7 +860,7 @@ absl::StatusOr EmitYnnSubgraph( continue; } - if (instr->IsElementwise()) { + if (instr->IsElementwise() && instr->opcode() != HloOpcode::kCopy) { if (!IsElementwiseOpSupportedByYnn(instr)) { return InvalidArgument( "Unsupported elementwise instruction in YNN fusion: %s", @@ -894,6 +909,16 @@ absl::StatusOr EmitYnnSubgraph( DefineReshapeOp(subgraph.get(), tensor_ids, instr)); } break; + case HloOpcode::kCopy: { + if (!IsCopyOpSupportedByYnn(instr)) { + return InvalidArgument( + "Unsupported copy instruction in YNN fusion: %s", + instr->ToString()); + } + ABSL_ASSIGN_OR_RETURN(tensor_ids[instr], + DefineCopyOp(subgraph.get(), tensor_ids, instr)); + } break; + case HloOpcode::kTranspose: { if (!IsTransposeOpSupportedByYnn(instr)) { return InvalidArgument( diff --git a/third_party/xla/xla/backends/cpu/ynn_support.cc b/third_party/xla/xla/backends/cpu/ynn_support.cc index 38bbfe050ab1ff..2ae056f728b000 100644 --- a/third_party/xla/xla/backends/cpu/ynn_support.cc +++ b/third_party/xla/xla/backends/cpu/ynn_support.cc @@ -17,7 +17,9 @@ limitations under the License. #include #include +#include #include +#include #include "ynnpack/include/ynnpack.h" #include "absl/algorithm/container.h" @@ -130,7 +132,16 @@ bool IsLayoutSupportedByYnn(const Shape& shape) { // TODO(b/460602165): We should eliminate this limitation. return false; } - return !shape.has_layout() || LayoutUtil::HasDescendingLayout(shape.layout()); + if (!shape.has_layout()) { + return true; + } + std::vector minor_to_major; + for (int64_t dim : shape.layout().minor_to_major()) { + if (shape.dimensions(dim) != 1) { + minor_to_major.push_back(dim); + } + } + return absl::c_is_sorted(minor_to_major, std::greater()); } namespace { @@ -165,6 +176,28 @@ bool IsBitcastOpSupportedByYnn(const HloInstruction* hlo) { return hlo->shape().element_type() == input->shape().element_type(); } +bool IsCopyOpSupportedByYnn(const HloInstruction* hlo) { + CHECK_EQ(hlo->opcode(), HloOpcode::kCopy); + if (!CheckOperandCount(hlo, 1)) { + return false; + } + if (!YnnType(hlo->shape().element_type()).ok()) { + return false; + } + const HloInstruction* input = hlo->operand(0); + if (hlo->shape().element_type() != input->shape().element_type()) { + return false; + } + if (hlo->shape().dimensions() != input->shape().dimensions()) { + return false; + } + if (!IsLayoutSupportedByYnn(hlo->shape()) || + !IsLayoutSupportedByYnn(input->shape())) { + return false; + } + return true; +} + bool IsReshapeOpSupportedByYnn(const HloInstruction* hlo) { CHECK_EQ(hlo->opcode(), HloOpcode::kReshape); if (!CheckOperandCount(hlo, 1)) { diff --git a/third_party/xla/xla/backends/cpu/ynn_support.h b/third_party/xla/xla/backends/cpu/ynn_support.h index de5ad5f2696eda..2eb3c3417d72ee 100644 --- a/third_party/xla/xla/backends/cpu/ynn_support.h +++ b/third_party/xla/xla/backends/cpu/ynn_support.h @@ -56,6 +56,9 @@ bool IsLayoutSupportedByYnn(const Shape& shape); // Returns true if the bitcast op is supported by YNNPACK. bool IsBitcastOpSupportedByYnn(const HloInstruction* hlo); +// Returns true if the copy op is supported by YNNPACK. +bool IsCopyOpSupportedByYnn(const HloInstruction* hlo); + // Returns true if the reshape op is supported by YNNPACK. bool IsReshapeOpSupportedByYnn(const HloInstruction* hlo); From 70e1a43ee8af29ed81f8ff03cf55ff47298ad68e Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Tue, 29 Sep 2026 10:23:23 -0700 Subject: [PATCH 24/25] Enable strict properties in assembly format for xtile dialect MLIR is migrating dialects to strict property handling in custom assembly formats (`useStrictPropertiesInAssemblyFormat = 1`), where `attr-dict` only parses and prints discardable attributes and inherent attributes must be either explicitly bound in the assembly format or handled via `prop-dict`. Add `prop-dict` before `attr-dict` in the assembly formats of `xtile.scan` and `xtile.dot_scaled`, enable `useStrictPropertiesInAssemblyFormat = 1` on `XTileDialect`, and update tests to use the `<...>` property dictionary syntax. PiperOrigin-RevId: 990387232 --- .../fusion_emitter_shared_dialect_test.cc | 6 ++--- .../tests/xtile_to_triton_lowering.mlir | 24 +++++++++---------- .../xla/xla/codegen/xtile/ir/tests/ops.mlir | 22 ++++++++--------- .../xla/xla/codegen/xtile/ir/xtile_dialect.td | 2 +- .../xla/xla/codegen/xtile/ir/xtile_ops.td | 4 ++-- 5 files changed, 29 insertions(+), 29 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc index 5653ba9c4e6013..4b12d7be11a5fe 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc @@ -292,7 +292,7 @@ ENTRY e { // CHECK: %[[EXTRACT1:.*]] = xtile.extract %arg1[] [] [] : memref -> tensor // CHECK: %[[INIT:.*]] = stablehlo.reshape %[[EXTRACT1]] : (tensor) -> tensor<1xf32> // CHECK: %[[OUTPUT:.*]], %{{.*}} = xtile.scan(%[[EXTRACT0]]) inits(%[[INIT]]) -// CHECK-SAME: dimension = 0 {scan_dim_size = 1024 : i64} +// CHECK-SAME: dimension = 0 // CHECK-SAME: : (tensor<1024xf32>), (tensor<1xf32>) -> (tensor<1024xf32>), (tensor<1xf32>) { // CHECK: ^bb0(%[[INPUT:.*]]: tensor, %[[CARRY:.*]]: tensor): // CHECK: %[[ADD:.*]] = stablehlo.add %[[INPUT]], %[[CARRY]] : tensor @@ -411,7 +411,7 @@ ENTRY e { EXPECT_OK(CreateXTileIrAndFileCheck( *module->GetComputationWithName("triton_dot"), block_level_parameters, R"( - CHECK: %[[DOT:.*]] = xtile.dot_scaled %[[LHS:.*]] scale %[[LHS_SCALE:.*]], %[[RHS:.*]] scale %[[RHS_SCALE:.*]] {dot_dimension_numbers = #stablehlo.dot, fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> + CHECK: %[[DOT:.*]] = xtile.dot_scaled %[[LHS:.*]] scale %[[LHS_SCALE:.*]], %[[RHS:.*]] scale %[[RHS_SCALE:.*]] > : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> CHECK: %[[RES:.*]] = arith.addf %{{.*}}, %[[DOT]] : tensor<128x256xf32> )")); } @@ -680,7 +680,7 @@ ENTRY e { CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[CARRY:.*]] = %{{.*}}) CHECK: %[[REVERSE_IDX:.*]] = arith.subi %{{.*}}, %[[IV]] : index CHECK: %[[INPUT_TILE:.*]] = xtile.extract %arg0 - CHECK: %[[SCAN_OUT:.*]], %[[NEW_CARRY:.*]] = xtile.scan(%[[INPUT_TILE]]) inits(%[[CARRY]]) dimension = 0 {is_reverse = true + CHECK: %[[SCAN_OUT:.*]], %[[NEW_CARRY:.*]] = xtile.scan(%[[INPUT_TILE]]) inits(%[[CARRY]]) dimension = 0 CHECK: xtile.insert %[[SCAN_OUT]] into %arg2 CHECK: scf.yield %[[NEW_CARRY]] )")); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir index 1657db56a13535..3bf693607595ff 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/xtile_to_triton_lowering.mlir @@ -26,8 +26,8 @@ func.func @lower_dot_scaled_add_to_triton( // CHECK: %[[RES:.*]] = tt.dot_scaled %[[LHS]] scale %[[LHS_SCALE]], %[[RHS]] scale %[[RHS_SCALE]], %[[ACC]] lhs = e5m2 rhs = e5m2 {fastMath = true} : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> // CHECK-NOT: arith.addf %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<128x128xf8E5M2>, + , + fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2> : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> %1 = arith.addf %acc, %0 : tensor<128x256xf32> // CHECK: return %[[RES]] : tensor<128x256xf32> @@ -44,8 +44,8 @@ func.func @lower_packed_dot_scaled_add_to_triton( // CHECK: %[[RES:.*]] = tt.dot_scaled %[[LHS]] scale %[[LHS_SCALE]], %[[RHS]] scale %[[RHS_SCALE]], %[[ACC]] lhs = e2m1 rhs = e2m1 {fastMath = true} : tensor<128x32xi8>, tensor<128x2xi8> * tensor<32x128xi8>, tensor<128x2xi8> -> tensor<128x128xf32> // CHECK-NOT: arith.addf %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f4E2M1FN, rhs_elem_type = f4E2M1FN} : tensor<128x32xi8>, + , + fastMath = true, lhs_elem_type = f4E2M1FN, rhs_elem_type = f4E2M1FN> : tensor<128x32xi8>, tensor<128x2xi8> * tensor<32x128xi8>, tensor<128x2xi8> -> tensor<128x128xf32> %1 = arith.addf %acc, %0 : tensor<128x128xf32> // CHECK: return %[[RES]] : tensor<128x128xf32> @@ -67,8 +67,8 @@ func.func @lower_dot_scaled_in_loop_non_canonical( // CHECK: %[[DOT:.*]] = tt.dot_scaled %{{.*}} scale %{{.*}}, %{{.*}} scale %{{.*}}, %[[ACC_2D]] // CHECK: scf.yield %[[DOT]] : tensor<128x256xf32> %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<1x128x128xf8E5M2>, tensor<1x128x4xi8> * tensor<1x128x256xf8E5M2>, tensor<1x256x4xi8> -> tensor<1x128x256xf32> + , + fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2> : tensor<1x128x128xf8E5M2>, tensor<1x128x4xi8> * tensor<1x128x256xf8E5M2>, tensor<1x256x4xi8> -> tensor<1x128x256xf32> %1 = arith.addf %accum, %0 : tensor<1x128x256xf32> scf.yield %1 : tensor<1x128x256xf32> } @@ -92,7 +92,7 @@ func.func @scan_lowering(%input0: tensor<16x16x16xf32>, %init0: tensor<16x16x1xf // CHECK-DAG: %[[RES1:.*]] = arith.addf %[[BCAST_INIT0]], %[[SCAN]] : tensor<16x16x16xf32> // CHECK: return %[[RES1]] : tensor<16x16x16xf32> - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -131,7 +131,7 @@ func.func @scan_lowering_carry(%input0: tensor<16x16x16xf32>, %init0: tensor<16x // CHECK-NOT: tt.reshape // CHECK: return %[[RES0]], %[[RES1]] - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 {scan_dim_size = 16 : i64} : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 2 : (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) -> (tensor<16x16x16xf32>), (tensor<16x16x1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -167,7 +167,7 @@ func.func @scan_lowering_1d_carry(%input0: tensor<16xf32>, %init0: tensor<1xf32> // CHECK-NOT: tt.gather // CHECK-NOT: tt.reshape // CHECK: return %[[RES0]], %[[RES1]] - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 {scan_dim_size = 16 : i64} : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -198,7 +198,7 @@ func.func @reverse_scan_lowering_carry(%input0: tensor<16xf32>, %init0: tensor<1 // CHECK-NOT: tt.gather // CHECK-NOT: tt.reshape // CHECK: return %[[RES0]], %[[RES1]] - %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 {scan_dim_size = 16 : i64, is_reverse = true} : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { + %0, %1 = xtile.scan(%input0) inits(%init0) dimension = 0 : (tensor<16xf32>), (tensor<1xf32>) -> (tensor<16xf32>), (tensor<1xf32>) { ^bb0(%arg0: f32, %arg1: f32): %add = arith.addf %arg0, %arg1 : f32 xtile.yield %add : f32 @@ -215,8 +215,8 @@ func.func @lower_dot_scaled_without_add_falls_back_to_xtile( -> tensor<128x256xf32> { // CHECK: %[[RES:.*]] = xtile.dot_scaled %[[LHS]] scale %[[LHS_SCALE]], %[[RHS]] scale %[[RHS_SCALE]] {{.*}}fastMath = true{{.*}} : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale - {dot_dimension_numbers = #stablehlo.dot, - fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2} : tensor<128x128xf8E5M2>, + , + fastMath = true, lhs_elem_type = f8E5M2, rhs_elem_type = f8E5M2> : tensor<128x128xf8E5M2>, tensor<128x4xi8> * tensor<128x256xf8E5M2>, tensor<256x4xi8> -> tensor<128x256xf32> // CHECK: return %[[RES]] : tensor<128x256xf32> return %0 : tensor<128x256xf32> diff --git a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir index d05c7e37e3741d..2f9e60aacc1f54 100644 --- a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir +++ b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir @@ -133,7 +133,7 @@ func.func @type_mismatch_insert(%src: tensor<24xf64>, %dst: memref<1024xf32>) { // ----- func.func @dot_scaled(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xi8>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>, %acc: tensor<128x256xf32>) -> tensor<128x256xf32> { - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = f32, rhs_elem_type = f32} : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -141,7 +141,7 @@ func.func @dot_scaled(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xi8>, % func.func @invalid_dot_scaled_scale_type(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xbf16>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>) -> tensor<128x256xf32> { // expected-error@+1 {{LHS scale element type 'bf16' is not supported}} - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = f32, rhs_elem_type = f32} : tensor<128x128xf32>, tensor<128x4xbf16> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xf32>, tensor<128x4xbf16> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -149,7 +149,7 @@ func.func @invalid_dot_scaled_scale_type(%lhs: tensor<128x128xf32>, %lhs_scale: func.func @invalid_dot_scaled_operand_type(%lhs: tensor<128x128xi32>, %lhs_scale: tensor<128x4xi8>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>) -> tensor<128x256xf32> { // expected-error@+1 {{'xtile.dot_scaled' op operand #0 must be ranked tensor of floating-point or 8-bit signless integer values, but got 'tensor<128x128xi32>'}} - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = f32, rhs_elem_type = f32} : tensor<128x128xi32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xi32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -157,7 +157,7 @@ func.func @invalid_dot_scaled_operand_type(%lhs: tensor<128x128xi32>, %lhs_scale func.func @invalid_dot_scaled_logical_elem_type(%lhs: tensor<128x128xf32>, %lhs_scale: tensor<128x4xi8>, %rhs: tensor<128x256xf32>, %rhs_scale: tensor<256x4xi8>) -> tensor<128x256xf32> { // expected-error@+1 {{LHS logical element type 'i32' is not supported}} - %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale {fastMath = true, lhs_elem_type = i32, rhs_elem_type = f32} : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> + %0 = xtile.dot_scaled %lhs scale %lhs_scale, %rhs scale %rhs_scale : tensor<128x128xf32>, tensor<128x4xi8> * tensor<128x256xf32>, tensor<256x4xi8> -> tensor<128x256xf32> return %0 : tensor<128x256xf32> } @@ -189,7 +189,7 @@ func.func @illegal_mask_out_of_bounds(%src: tensor<32xf64>, %mask: f64) -> tenso // ----- func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -201,7 +201,7 @@ func.func @scan_op(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> func.func @scan_op_rank_reduced_init(%input: tensor<10x20x30xf32>, %init: tensor<10x20xf32>) -> (tensor<10x20x30xf32>, tensor<10x20xf32>) { // expected-error @+1 {{init #0 type 'tensor<10x20xf32>' does not match expected type 'tensor<10x20x1xf32>'}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -213,7 +213,7 @@ func.func @scan_op_rank_reduced_init(%input: tensor<10x20x30xf32>, %init: tensor func.func @scan_op_dimension_out_of_range(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { // expected-error @+1 {{scan dimension 3 is out of range for input #0 of rank 3}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 3 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 3 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -225,7 +225,7 @@ func.func @scan_op_dimension_out_of_range(%input: tensor<10x20x30xf32>, %init: t func.func @scan_op_init_count_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { // expected-error @+1 {{expects one init, output and carry per input, but got 1 inputs, 2 inits, 1 outputs and 1 carries}} - %output, %carry = xtile.scan(%input) inits(%init, %init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>, tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init, %init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>, tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -237,7 +237,7 @@ func.func @scan_op_init_count_mismatch(%input: tensor<10x20x30xf32>, %init: tens func.func @scan_op_negative_dimension(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x1xf32>) { // expected-error @+1 {{scan dimension -1 is out of range for input #0 of rank 3}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = -1 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = -1 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -249,7 +249,7 @@ func.func @scan_op_negative_dimension(%input: tensor<10x20x30xf32>, %init: tenso func.func @scan_op_output_type_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>, tensor<10x20x1xf32>) { // expected-error @+1 {{output #0 type 'tensor<10x20x30xf64>' does not match input type 'tensor<10x20x30xf32>'}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>), (tensor<10x20x1xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf64>), (tensor<10x20x1xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> @@ -261,7 +261,7 @@ func.func @scan_op_output_type_mismatch(%input: tensor<10x20x30xf32>, %init: ten func.func @scan_op_carry_type_mismatch(%input: tensor<10x20x30xf32>, %init: tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>, tensor<10x20x30xf32>) { // expected-error @+1 {{carry #0 type 'tensor<10x20x30xf32>' does not match expected type 'tensor<10x20x1xf32>'}} - %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 {scan_dim_size = 30 : i64} : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x30xf32>) { + %output, %carry = xtile.scan(%input) inits(%init) dimension = 2 : (tensor<10x20x30xf32>), (tensor<10x20x1xf32>) -> (tensor<10x20x30xf32>), (tensor<10x20x30xf32>) { ^bb0(%arg0: tensor<10x20xf32>, %arg1: tensor<10x20xf32>): %add = stablehlo.add %arg0, %arg1 : tensor<10x20xf32> stablehlo.return %add : tensor<10x20xf32> diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td index 8a9d655d3b71cc..38c4bbf47a1ce8 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td @@ -34,7 +34,7 @@ def XTileDialect : Dialect { ]; let useDefaultAttributePrinterParser = 1; - let useStrictPropertiesInAssemblyFormat = 0; + let useStrictPropertiesInAssemblyFormat = 1; } #endif // XLA_CODEGEN_XTILE_IR_XTILE_DIALECT diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td index 1cbaab2442e73b..468a9b05823561 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td @@ -350,7 +350,7 @@ def ScanOp : XTile_Op<"scan", [Pure, AttrSizedOperandSegments, AttrSizedResultSe let assemblyFormat = [{ `(` $inputs `)` `inits` `(` $inits `)` `dimension` `=` $dimension - attr-dict + prop-dict attr-dict `:` `(` type($inputs) `)` `,` `(` type($inits) `)` `->` `(` type($outputs) `)` `,` `(` type($carries) `)` $body }]; @@ -391,7 +391,7 @@ def DotScaledOp : XTile_Op<"dot_scaled", [Pure, AttrSizedOperandSegments]> { let results = (outs RankedTensorOf<[AnyFloat]>:$result); let assemblyFormat = [{ - $lhs (`scale` $lhs_scale^)? `,` $rhs (`scale` $rhs_scale^)? attr-dict + $lhs (`scale` $lhs_scale^)? `,` $rhs (`scale` $rhs_scale^)? prop-dict attr-dict `:` type($lhs) (`,` type($lhs_scale)^)? `*` type($rhs) (`,` type($rhs_scale)^)? `->` type($result) }]; From ac3addad8afc055a23f362203581cd362d92c3ea Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Tue, 29 Sep 2026 10:46:40 -0700 Subject: [PATCH 25/25] [XLA] Do not give one sided intervals a symmetric magnitude bound in the constraint propagator. PiperOrigin-RevId: 990402570 --- .../xla/xla/tests/constraint_propagator.cc | 2 +- .../xla/tests/constraint_propagator_test.cc | 18 ++++++++---------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index efc9489b1e2602..ca487d25a098a7 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -329,7 +329,7 @@ std::optional GetSymmetricMagnitudeBound( const ConstraintInterval& interval) { if ((interval.min == ConstraintInterval::kMin && interval.max == ConstraintInterval::kMax) || - interval.min > 0.0 || interval.max < 0.0) { + interval.min >= 0.0 || interval.max <= 0.0) { return std::nullopt; } double m = std::min(std::abs(interval.min), std::abs(interval.max)); diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index cd35a5753dd931..e7f9f1e32677cf 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -157,10 +157,8 @@ ENTRY main { auto y_int = states[module->entry_computation()->parameter_instruction(1)] .GetConstraintInterval(); - EXPECT_TRUE(x_int.IsPositive()); - EXPECT_FALSE(x_int.exclude_zero); - EXPECT_TRUE(y_int.IsNegative()); - EXPECT_FALSE(y_int.exclude_zero); + EXPECT_EQ(x_int, ConstraintInterval::Positive()); + EXPECT_EQ(y_int, ConstraintInterval::Negative()); } TEST_F(ConstraintPropagatorTest, SubtractNegative) { @@ -183,10 +181,8 @@ ENTRY main { auto y_int = states[module->entry_computation()->parameter_instruction(1)] .GetConstraintInterval(); - EXPECT_TRUE(x_int.IsNegative()); - EXPECT_FALSE(x_int.exclude_zero); - EXPECT_TRUE(y_int.IsPositive()); - EXPECT_FALSE(y_int.exclude_zero); + EXPECT_EQ(x_int, ConstraintInterval::Negative()); + EXPECT_EQ(y_int, ConstraintInterval::Positive()); } TEST_F(ConstraintPropagatorTest, RsqrtMulForcedNegative) { @@ -961,9 +957,11 @@ ENTRY main { auto p1_int = states[module->entry_computation()->parameter_instruction(1)] .GetConstraintInterval(); EXPECT_FALSE(p0_int.IsEmpty()); - EXPECT_LE(p0_int.max, 0.0); + EXPECT_DOUBLE_EQ(p0_int.min, -4.0); + EXPECT_DOUBLE_EQ(p0_int.max, 0.0); EXPECT_FALSE(p1_int.IsEmpty()); - EXPECT_GE(p1_int.min, 0.0); + EXPECT_DOUBLE_EQ(p1_int.min, 0.0); + EXPECT_DOUBLE_EQ(p1_int.max, 4.0); } // Tests that for Power(base, exp) with base in (0, 1) (decay schedule, e.g.