From 87372d77db3f353cb36fbf82ed4503073e565895 Mon Sep 17 00:00:00 2001 From: Crefeda Rodrigues Date: Mon, 7 Sep 2026 13:16:43 +0000 Subject: [PATCH 01/25] Add AArch64 oneDNN async module repo --- MODULE.bazel | 1 + 1 file changed, 1 insertion(+) diff --git a/MODULE.bazel b/MODULE.bazel index 539d8107d2fa2f..5105fe8c311d61 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -133,6 +133,7 @@ use_repo( "nvshmem", "onednn", "onednn_async", + "onednn_async_aarch64", "pthreadpool", "raft", "riegeli", From b599ff6ebd4f782a801ccee2efa7cc4bf325ebc3 Mon Sep 17 00:00:00 2001 From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:43:37 -0700 Subject: [PATCH 02/25] Validate axis bounds in experimental.numpy reduction functions NumPy raises AxisError when sum/prod/mean/amax/amin/std/var receive an out-of-bounds axis. The shared _reduce helper forwarded the axis to the backend kernel unvalidated, producing a confusing runtime failure. Add static bounds validation for sequence axes with statically known rank, consistent with the cumsum/compress/moveaxis validation pattern. --- tensorflow/python/ops/numpy_ops/np_array_ops.py | 16 ++++++++++++++++ .../python/ops/numpy_ops/np_array_ops_test.py | 11 +++++++++++ 2 files changed, 27 insertions(+) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 250893b1e4797c..a1c7bbf285b996 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -620,6 +620,22 @@ def _reduce( if keepdims is None: keepdims = False a = asarray(a, dtype=dtype) + # NumPy raises AxisError for out-of-bounds axes instead of letting the + # backend kernel fail with a confusing error. + maybe_rank = a.shape.rank + if ( + maybe_rank is not None + and axis is not None + and isinstance(axis, (list, tuple, range, np.ndarray)) + and builtins.all(isinstance(x, (int, np.integer)) for x in axis) + ): + for ax in axis: + normalized = ax + maybe_rank if ax < 0 else ax + if normalized < 0 or normalized >= maybe_rank: + raise ValueError( + f'Argument `axis` (received axis={ax}) is out of bounds ' + f'for input of rank {maybe_rank}.' + ) if ( dtype == np.bool_ or preserve_bool and a.dtype == np.bool_ ) and tf_bool_fn is not None: diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py index 6ba00dfc708568..24e338cd961155 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py @@ -1022,9 +1022,20 @@ def run_test(a, **kwargs): run_test([[2, -3], [-6, 7]], axis=1, keepdims=True) run_test([[2, -3], [-6, 7]], axis=(0, 1)) run_test([[2, -3], [-6, 7]], axis=(1, 0)) + with self.assertRaisesRegex(ValueError, 'out of bounds'): + math_fun(np_array_ops.array([[2, -3], [-6, 7]]), axis=2) + with self.assertRaisesRegex(ValueError, 'out of bounds'): + math_fun(np_array_ops.array([[2, -3], [-6, 7]]), axis=-3) + with self.assertRaisesRegex(ValueError, 'out of bounds'): + math_fun(np_array_ops.array([[2, -3], [-6, 7]]), axis=(0, 2)) def testSum(self): self._testReduce(np_array_ops.sum, np.sum, 'sum') + # Scalar input: axes -1 and 0 are in bounds, 1 and -2 are not. + with self.assertRaisesRegex(ValueError, 'out of bounds'): + np_array_ops.sum(5, axis=1) + with self.assertRaisesRegex(ValueError, 'out of bounds'): + np_array_ops.sum(5, axis=-2) def testAmax(self): self._testReduce(np_array_ops.amax, np.amax, 'amax') From bb4d7e2b614868c72494cee9e210e083b5b924e8 Mon Sep 17 00:00:00 2001 From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:47:04 -0700 Subject: [PATCH 03/25] Handle scalar and 0-d array axes in reduction axis validation --- .../python/ops/numpy_ops/np_array_ops.py | 34 ++++++++++++------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index a1c7bbf285b996..e498d2801b51f6 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -623,19 +623,27 @@ def _reduce( # NumPy raises AxisError for out-of-bounds axes instead of letting the # backend kernel fail with a confusing error. maybe_rank = a.shape.rank - if ( - maybe_rank is not None - and axis is not None - and isinstance(axis, (list, tuple, range, np.ndarray)) - and builtins.all(isinstance(x, (int, np.integer)) for x in axis) - ): - for ax in axis: - normalized = ax + maybe_rank if ax < 0 else ax - if normalized < 0 or normalized >= maybe_rank: - raise ValueError( - f'Argument `axis` (received axis={ax}) is out of bounds ' - f'for input of rank {maybe_rank}.' - ) + if maybe_rank is not None and axis is not None: + # Wrap scalar axes into a sequence so both ints and sequences of ints + # are validated; 0-d NumPy arrays cannot be iterated directly. + if isinstance(axis, (int, np.integer)): + static_axes = (axis,) + elif isinstance(axis, np.ndarray) and axis.ndim == 0: + static_axes = (int(axis),) + elif isinstance(axis, (list, tuple, range, np.ndarray)): + static_axes = axis + else: + static_axes = None + if static_axes is not None and builtins.all( + isinstance(ax, (int, np.integer)) for ax in static_axes + ): + for ax in static_axes: + normalized = ax + maybe_rank if ax < 0 else ax + if normalized < 0 or normalized >= maybe_rank: + raise ValueError( + f'Argument `axis` (received axis={ax}) is out of bounds ' + f'for input of rank {maybe_rank}.' + ) if ( dtype == np.bool_ or preserve_bool and a.dtype == np.bool_ ) and tf_bool_fn is not None: From f137c3e51740943eac52ccf18c7d440fa85423c1 Mon Sep 17 00:00:00 2001 From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com> Date: Mon, 21 Sep 2026 18:33:28 -0700 Subject: [PATCH 04/25] Expect ValueError for out-of-bounds reduction axis in testReductionOfOutOfBoundsAxis --- tensorflow/python/ops/numpy_ops/tests/np_test.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py index 054549a78fea00..0a87596422256e 100644 --- a/tensorflow/python/ops/numpy_ops/tests/np_test.py +++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py @@ -27,7 +27,6 @@ import six from tensorflow.python.framework import errors -from tensorflow.python.framework import errors_impl from tensorflow.python.ops.numpy_ops.tests.config import config from tensorflow.python.ops.numpy_ops.tests.config import FLAGS import tensorflow.python.ops.numpy_ops.tests.extensions as nje @@ -39,7 +38,9 @@ config.parse_flags_with_absl() -nonempty_nonscalar_array_shapes = [(4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)] +nonempty_nonscalar_array_shapes = [ + (4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4) +] nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes empty_array_shapes = [(0,), (0, 4), (3, 0),] @@ -2601,9 +2602,7 @@ def f(x, v): def testReductionOfOutOfBoundsAxis(self): # Issue 888 x = tnp.ones((3, 4)) - self.assertRaises( - errors_impl.InvalidArgumentError, lambda: tnp.sum(x, axis=2) - ) + self.assertRaises(ValueError, lambda: tnp.sum(x, axis=2)) @jtu.disable def testIssue956(self): From a712eb48e95e69cc59d3e3c8cce944b1648861b7 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 22 Sep 2026 16:58:00 -0400 Subject: [PATCH 05/25] Reject negative integer powers on GPU --- tensorflow/core/kernels/BUILD | 7 + tensorflow/core/kernels/cwise_op_pow.cc | 7 +- tensorflow/core/kernels/cwise_ops.h | 5 +- tensorflow/core/kernels/integer_pow_gpu.cu.cc | 57 ++++++++ tensorflow/core/kernels/integer_pow_gpu.h | 125 ++++++++++++++++++ tensorflow/core/kernels/mlir_generated/BUILD | 1 + .../core/kernels/mlir_generated/gpu_op_pow.cc | 28 +++- .../math_ops/cwise_ops_binary_test.py | 36 ++++- 8 files changed, 252 insertions(+), 14 deletions(-) create mode 100644 tensorflow/core/kernels/integer_pow_gpu.cu.cc create mode 100644 tensorflow/core/kernels/integer_pow_gpu.h diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index dbc1438ef90514..a978c30bb21ea7 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -3816,6 +3816,12 @@ tf_kernel_library( deps = MATH_DEPS + ["//tensorflow/core:framework_internal"], ) +tf_kernel_library( + name = "integer_pow_gpu", + prefix = "integer_pow_gpu", + deps = MATH_DEPS + ["//tensorflow/core:framework_internal"], +) + tf_kernel_library( name = "cross_op", prefix = "cross_op", @@ -3837,6 +3843,7 @@ tf_kernel_library( ), prefix = "cwise_op", deps = MATH_DEPS + [ + ":integer_pow_gpu", "//tensorflow/core/kernels/mlir_generated:cwise_op", "@com_google_absl//absl/base:prefetch", "@com_google_absl//absl/status", diff --git a/tensorflow/core/kernels/cwise_op_pow.cc b/tensorflow/core/kernels/cwise_op_pow.cc index ae21c4613f1bc4..4a3d07c501511a 100644 --- a/tensorflow/core/kernels/cwise_op_pow.cc +++ b/tensorflow/core/kernels/cwise_op_pow.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/core/kernels/cwise_ops_common.h" +#include "tensorflow/core/kernels/integer_pow_gpu.h" namespace tensorflow { REGISTER6(BinaryOp, CPU, "Pow", functor::pow, float, Eigen::half, bfloat16, @@ -24,7 +25,11 @@ REGISTER4(BinaryOp, CPU, "Pow", functor::safe_pow, int8_t, int16_t, int32_t, #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) REGISTER3(BinaryOp, GPU, "Pow", functor::pow, float, Eigen::half, double); -REGISTER(BinaryOp, GPU, "Pow", functor::safe_pow_ignore_error, int64); +using Int64GpuPowOp = + IntegerPowGpuOp>>; +REGISTER_KERNEL_BUILDER( + Name("Pow").Device(DEVICE_GPU).TypeConstraint("T"), Int64GpuPowOp); #endif REGISTER(BinaryOp, GPU, "Pow", functor::pow, bfloat16); #endif diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h index 7c0c3f88fe4a74..236cb5711bd18a 100644 --- a/tensorflow/core/kernels/cwise_ops.h +++ b/tensorflow/core/kernels/cwise_ops.h @@ -1179,8 +1179,9 @@ struct safe_pow : base> { static constexpr bool has_errors = true; }; -// Version of safe_pow for integers which returns 0 if RHS is negative and LHS -// is not 1 or -1. For use on GPUs, where we cannot raise an error. +// Version of safe_pow for GPU computation. The integer Pow GPU kernel checks +// negative exponents separately and reports an error after the device work +// completes; this functor must still produce a defined result on the device. template struct safe_pow_ignore_error_op { static_assert(std::is_integral::value, "Integer type expected"); diff --git a/tensorflow/core/kernels/integer_pow_gpu.cu.cc b/tensorflow/core/kernels/integer_pow_gpu.cu.cc new file mode 100644 index 00000000000000..78922c1812c6b7 --- /dev/null +++ b/tensorflow/core/kernels/integer_pow_gpu.cu.cc @@ -0,0 +1,57 @@ +/* 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. +==============================================================================*/ + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + +#define EIGEN_USE_GPU +#include "tensorflow/core/kernels/integer_pow_gpu.h" + +#include +#include + +#include "tensorflow/core/util/gpu_kernel_helper.h" + +namespace tensorflow { +namespace { + +template +__global__ void CheckNegativeExponentKernel(const T* exponents, int64_t size, + int32_t* found_negative) { + for (int64_t i : GpuGridRangeX(size)) { + if (exponents[i] < 0) { + atomicExch(found_negative, 1); + } + } +} + +} // namespace + +template +absl::Status CheckNegativeExponentGpu::operator()( + const Eigen::GpuDevice& device, const T* exponents, int64_t size, + int32_t* found_negative) { + constexpr int kThreads = 256; + const int blocks = std::min(1 + (size - 1) / kThreads, 1024); + return GpuLaunchKernel(CheckNegativeExponentKernel, blocks, kThreads, 0, + device.stream(), exponents, size, found_negative); +} + +template struct CheckNegativeExponentGpu; +template struct CheckNegativeExponentGpu; +template struct CheckNegativeExponentGpu; + +} // namespace tensorflow + +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/tensorflow/core/kernels/integer_pow_gpu.h b/tensorflow/core/kernels/integer_pow_gpu.h new file mode 100644 index 00000000000000..66e979901dce18 --- /dev/null +++ b/tensorflow/core/kernels/integer_pow_gpu.h @@ -0,0 +1,125 @@ +/* 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. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_KERNELS_INTEGER_POW_GPU_H_ +#define TENSORFLOW_CORE_KERNELS_INTEGER_POW_GPU_H_ + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + +#include +#include +#include + +#include "absl/status/status.h" +#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/tensor_reference.h" +#include "tensorflow/core/framework/tensor_types.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/util/bcast.h" + +namespace tensorflow { + +// Instantiated in integer_pow_gpu.cu.cc for the signed GPU Pow types. +template +struct CheckNegativeExponentGpu { + absl::Status operator()(const Eigen::GpuDevice& device, const T* exponents, + int64_t size, int32_t* found_negative); +}; + +// Applies the same exponent validation to the Eigen and MLIR GPU kernels. +// The check runs before Pow can reuse an exponent buffer for its output. Only +// a single flag is copied back to the host, without blocking the caller. +template +class IntegerPowGpuOp : public AsyncOpKernel { + public: + explicit IntegerPowGpuOp(OpKernelConstruction* context) + : AsyncOpKernel(context), kernel_(context) {} + + void ComputeAsync(OpKernelContext* context, DoneCallback done) override { + const BCast bcast(BCast::FromShape(context->input(0).shape()), + BCast::FromShape(context->input(1).shape())); + if (!bcast.IsValid() || + BCast::ToShape(bcast.output_shape()).num_elements() == 0) { + kernel_.Compute(context); + done(); + return; + } + + auto* stream = context->op_device_context()->stream(); + OP_REQUIRES_ASYNC(context, stream != nullptr, + absl::InternalError("No GPU stream available."), done); + + Tensor found_negative; + OP_REQUIRES_OK_ASYNC( + context, + context->allocate_temp(DT_INT32, TensorShape({}), &found_negative), + done); + stream_executor::DeviceAddressBase device_flag( + found_negative.flat().data(), sizeof(int32_t)); + OP_REQUIRES_OK_ASYNC( + context, stream->Memset32(&device_flag, 0, sizeof(int32_t)), done); + const auto& exponents = context->input(1).flat(); + OP_REQUIRES_OK_ASYNC( + context, + CheckNegativeExponentGpu()(context->eigen_device(), + exponents.data(), exponents.size(), + found_negative.flat().data()), + done); + + AllocatorAttributes attributes; + attributes.set_on_host(true); + attributes.set_gpu_compatible(true); + Tensor host_flag; + OP_REQUIRES_OK_ASYNC(context, + context->allocate_temp(DT_INT32, TensorShape({}), + &host_flag, attributes), + done); + OP_REQUIRES_ASYNC( + context, + stream + ->Memcpy(host_flag.flat().data(), device_flag, + sizeof(int32_t)) + .ok(), + absl::InternalError("GPU memcpy from device to host failed"), done); + + kernel_.Compute(context); + + TensorReference flag_ref(found_negative); + auto check = [stream, flag_ref, host_flag, context, done]() { + { + std::unique_ptr activation = + stream->parent()->Activate(); + flag_ref.Unref(); + if (context->status().ok() && host_flag.scalar()() != 0) { + context->SetStatus(absl::InvalidArgumentError( + "Integers to negative integer powers are not allowed")); + } + } + done(); + }; + context->device() + ->tensorflow_accelerator_device_info() + ->event_mgr->ThenExecute(stream, std::move(check)); + } + + private: + Kernel kernel_; +}; + +} // namespace tensorflow + +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#endif // TENSORFLOW_CORE_KERNELS_INTEGER_POW_GPU_H_ diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD index 1613f53ff7670d..0e7075f9c54a57 100644 --- a/tensorflow/core/kernels/mlir_generated/BUILD +++ b/tensorflow/core/kernels/mlir_generated/BUILD @@ -234,6 +234,7 @@ tf_kernel_library( tags = ["manual"], deps = if_mlir_generated_gpu_kernels_enabled([ ":base_gpu_op", + "//tensorflow/core/kernels:integer_pow_gpu", ":gpu_atan2_kernels", ":gpu_bitwise_and_kernels", ":gpu_bitwise_or_kernels", diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc index 256943985b5200..0f359808377502 100644 --- a/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc +++ b/tensorflow/core/kernels/mlir_generated/gpu_op_pow.cc @@ -12,19 +12,39 @@ 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 "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/kernels/integer_pow_gpu.h" #include "tensorflow/core/kernels/mlir_generated/base_gpu_op.h" +#include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive namespace tensorflow { GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_HALF); GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_FLOAT); GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_DOUBLE); -GENERATE_AND_REGISTER_BINARY_GPU_KERNEL(Pow, DT_INT64); +GENERATE_BINARY_GPU_KERNEL(Pow, DT_INT64); +using Int64GpuPowOp = + IntegerPowGpuOp; +REGISTER_KERNEL_BUILDER( + Name("Pow").Device(DEVICE_GPU).TypeConstraint("T"), Int64GpuPowOp); // These kernels are JIT-compiled. -GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Pow, DT_INT8); -GENERATE_AND_REGISTER_BINARY_JIT_GPU_KERNEL(Pow, DT_INT16); +GENERATE_BINARY_GPU_KERNEL(Pow, DT_INT8); +using Int8GpuPowOp = + IntegerPowGpuOp; +REGISTER_KERNEL_BUILDER(Name("Pow") + .Device(DEVICE_GPU) + .TypeConstraint("T") + .Label(kJitKernelLabel), + Int8GpuPowOp); + +GENERATE_BINARY_GPU_KERNEL(Pow, DT_INT16); +using Int16GpuPowOp = + IntegerPowGpuOp; +REGISTER_KERNEL_BUILDER(Name("Pow") + .Device(DEVICE_GPU) + .TypeConstraint("T") + .Label(kJitKernelLabel), + Int16GpuPowOp); } // namespace tensorflow diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py index b0cf1d0058c1a5..cef2b082165599 100644 --- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py +++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py @@ -848,7 +848,7 @@ def testAtan2SpecialValues(self): self._compareGpu(x1, x2, np.arctan2, math_ops.atan2) def testPowNegativeExponentCpu(self): - for dtype in [np.int32, np.int64]: + for dtype in [np.int8, np.int16, np.int32, np.int64]: with test_util.force_cpu(): with self.assertRaisesRegex( errors_impl.InvalidArgumentError, @@ -876,12 +876,34 @@ def testPowNegativeExponentCpu(self): def testPowNegativeExponentGpu(self): if not test_util.is_gpu_available(): self.skipTest("Requires GPU") - # Negative integer powers return zero on GPUs for abs(LHS) > 1. Negative - # integer powers for 1 and -1 will return the correct result. - x = np.array([2, 3, 1, -1, -1]).astype(np.int64) - y = np.array([-1, 0, -2, -2, -3]).astype(np.int64) - z = math_ops.pow(x, y) - self.assertAllEqual(self.evaluate(z), [0, 1, 1, 1, -1]) + for dtype in [np.int8, np.int16, np.int64]: + x = np.array([2, 3, 1, -1, -1], dtype=dtype) + y = np.array([-1, 0, -2, -2, -3], dtype=dtype) + with test_util.force_gpu(): + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + "Integers to negative integer powers are not allowed"): + self.evaluate(math_ops.pow(x, y)) + + # Check both scalar and broadcasted exponents, including -1 bases. + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + "Integers to negative integer powers are not allowed"): + self.evaluate(math_ops.pow(np.array([-1, 1], dtype=dtype), + np.array(-1, dtype=dtype))) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, + "Integers to negative integer powers are not allowed"): + self.evaluate(math_ops.pow(np.array([[2], [3]], dtype=dtype), + np.array([[1, -1]], dtype=dtype))) + self.assertAllEqual( + self.evaluate(math_ops.pow(np.array([[2], [3]], dtype=dtype), + np.array([[0, 2]], dtype=dtype))), + [[1, 4], [1, 9]]) + self.assertAllEqual( + self.evaluate(math_ops.pow(np.empty((0, 1), dtype=dtype), + np.array([[1, -1]], dtype=dtype))), + np.empty((0, 2), dtype=dtype)) @test.disable_with_predicate( pred=test.is_built_with_rocm, skip_message="On ROCm this test fails" From 094f260f704fce33d178f347aaed9a01caa1d546 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 22 Sep 2026 17:01:39 -0400 Subject: [PATCH 06/25] Simplify GPU Pow validation callback --- tensorflow/core/kernels/integer_pow_gpu.cu.cc | 3 ++- tensorflow/core/kernels/integer_pow_gpu.h | 15 +++++---------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/tensorflow/core/kernels/integer_pow_gpu.cu.cc b/tensorflow/core/kernels/integer_pow_gpu.cu.cc index 78922c1812c6b7..9549457fb40be3 100644 --- a/tensorflow/core/kernels/integer_pow_gpu.cu.cc +++ b/tensorflow/core/kernels/integer_pow_gpu.cu.cc @@ -43,7 +43,8 @@ absl::Status CheckNegativeExponentGpu::operator()( const Eigen::GpuDevice& device, const T* exponents, int64_t size, int32_t* found_negative) { constexpr int kThreads = 256; - const int blocks = std::min(1 + (size - 1) / kThreads, 1024); + const int blocks = static_cast( + std::min((size + kThreads - 1) / kThreads, 1024)); return GpuLaunchKernel(CheckNegativeExponentKernel, blocks, kThreads, 0, device.stream(), exponents, size, found_negative); } diff --git a/tensorflow/core/kernels/integer_pow_gpu.h b/tensorflow/core/kernels/integer_pow_gpu.h index 66e979901dce18..123fdf996ce81b 100644 --- a/tensorflow/core/kernels/integer_pow_gpu.h +++ b/tensorflow/core/kernels/integer_pow_gpu.h @@ -19,7 +19,6 @@ limitations under the License. #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #include -#include #include #include "absl/status/status.h" @@ -98,15 +97,11 @@ class IntegerPowGpuOp : public AsyncOpKernel { kernel_.Compute(context); TensorReference flag_ref(found_negative); - auto check = [stream, flag_ref, host_flag, context, done]() { - { - std::unique_ptr activation = - stream->parent()->Activate(); - flag_ref.Unref(); - if (context->status().ok() && host_flag.scalar()() != 0) { - context->SetStatus(absl::InvalidArgumentError( - "Integers to negative integer powers are not allowed")); - } + auto check = [flag_ref, host_flag, context, done]() { + flag_ref.Unref(); + if (context->status().ok() && host_flag.scalar()() != 0) { + context->SetStatus(absl::InvalidArgumentError( + "Integers to negative integer powers are not allowed")); } done(); }; From 6a1d2367842ca5dd57ff07c0cdd69958b9fe7af2 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Tue, 22 Sep 2026 17:11:23 -0400 Subject: [PATCH 07/25] Retain GPU validation tensor safely --- tensorflow/core/kernels/cwise_op_pow.cc | 7 +++---- tensorflow/core/kernels/integer_pow_gpu.h | 7 +++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/kernels/cwise_op_pow.cc b/tensorflow/core/kernels/cwise_op_pow.cc index 4a3d07c501511a..73edeae5cfd520 100644 --- a/tensorflow/core/kernels/cwise_op_pow.cc +++ b/tensorflow/core/kernels/cwise_op_pow.cc @@ -25,11 +25,10 @@ REGISTER4(BinaryOp, CPU, "Pow", functor::safe_pow, int8_t, int16_t, int32_t, #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #if !defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) REGISTER3(BinaryOp, GPU, "Pow", functor::pow, float, Eigen::half, double); -using Int64GpuPowOp = - IntegerPowGpuOp>>; +using Int64GpuPowOp = IntegerPowGpuOp< + int64_t, BinaryOp>>; REGISTER_KERNEL_BUILDER( - Name("Pow").Device(DEVICE_GPU).TypeConstraint("T"), Int64GpuPowOp); + Name("Pow").Device(DEVICE_GPU).TypeConstraint("T"), Int64GpuPowOp); #endif REGISTER(BinaryOp, GPU, "Pow", functor::pow, bfloat16); #endif diff --git a/tensorflow/core/kernels/integer_pow_gpu.h b/tensorflow/core/kernels/integer_pow_gpu.h index 123fdf996ce81b..65f15364e5173f 100644 --- a/tensorflow/core/kernels/integer_pow_gpu.h +++ b/tensorflow/core/kernels/integer_pow_gpu.h @@ -24,7 +24,6 @@ limitations under the License. #include "absl/status/status.h" #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" #include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/tensor_reference.h" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/platform/errors.h" #include "tensorflow/core/util/bcast.h" @@ -96,9 +95,9 @@ class IntegerPowGpuOp : public AsyncOpKernel { kernel_.Compute(context); - TensorReference flag_ref(found_negative); - auto check = [flag_ref, host_flag, context, done]() { - flag_ref.Unref(); + auto check = [found_negative, host_flag, context, done]() { + // Keep the device flag alive until the queued GPU work has completed. + (void)found_negative; if (context->status().ok() && host_flag.scalar()() != 0) { context->SetStatus(absl::InvalidArgumentError( "Integers to negative integer powers are not allowed")); From 05f42d29968a5f44514218cba04544bb5967ca37 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Wed, 23 Sep 2026 09:25:44 -0400 Subject: [PATCH 08/25] Use valid integer exponents in generated GPU Pow tests --- .../kernels/mlir_generated/gpu_binary_ops_test.cc | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc index 51087da47053c3..b1f041d562fe51 100644 --- a/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc +++ b/tensorflow/core/kernels/mlir_generated/gpu_binary_ops_test.cc @@ -1304,6 +1304,13 @@ absl::InlinedVector PowInput() { return test::InputAsVector({-2, -1, -1, 1, 1, 3}); } +template +absl::InlinedVector IntegerPowExponents() { + // These generated tests expect successful results. Negative exponents are + // covered separately by the GPU Pow error tests. + return test::InputAsVector({0, 1, 2, 1, 2, 3}); +} + template <> Eigen::half baseline_pow(Eigen::half lhs, Eigen::half rhs) { return static_cast( @@ -1327,18 +1334,18 @@ GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES( GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES( Pow, /*test_name=*/Int64, int64_t, int64_t, PowInput(), - PowInput(), baseline_pow, + IntegerPowExponents(), baseline_pow, test::OpsTestConfig().ExpectStrictlyEqual()) /// Test the JIT-compiled kernels. #if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES( Pow, /*test_name=*/Int8, int8_t, int8_t, PowInput(), - PowInput(), baseline_pow, + IntegerPowExponents(), baseline_pow, test::OpsTestConfig().ExpectStrictlyEqual()) GENERATE_DEFAULT_TESTS_WITH_SPECIFIC_INPUT_VALUES( Pow, /*test_name=*/Int16, int16_t, int16_t, PowInput(), - PowInput(), baseline_pow, + IntegerPowExponents(), baseline_pow, test::OpsTestConfig().ExpectStrictlyEqual()) #endif From 8084b92b3d059127191bbebfc23c2709799f9733 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Thu, 24 Sep 2026 06:38:40 -0400 Subject: [PATCH 09/25] Declare Abseil status dependency for GPU integer Pow --- tensorflow/core/kernels/BUILD | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index a978c30bb21ea7..226a6f3d4b84f7 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -3819,7 +3819,10 @@ tf_kernel_library( tf_kernel_library( name = "integer_pow_gpu", prefix = "integer_pow_gpu", - deps = MATH_DEPS + ["//tensorflow/core:framework_internal"], + deps = MATH_DEPS + [ + "//tensorflow/core:framework_internal", + "@com_google_absl//absl/status", + ], ) tf_kernel_library( From 1e52bce08f8449a571f2edc2de8828bf9bf12775 Mon Sep 17 00:00:00 2001 From: Bhushan Asati Date: Thu, 24 Sep 2026 21:37:13 -0400 Subject: [PATCH 10/25] Account for input aliasing in forward-mode function caches --- tensorflow/python/eager/forwardprop_test.py | 75 +++++++++++++++++++ .../polymorphic_function/concrete_function.py | 12 ++- 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/eager/forwardprop_test.py b/tensorflow/python/eager/forwardprop_test.py index 7fedd4eb558cf3..f87f4aa30224b1 100644 --- a/tensorflow/python/eager/forwardprop_test.py +++ b/tensorflow/python/eager/forwardprop_test.py @@ -657,6 +657,81 @@ def f(x): self.assertAllClose(3.5 * 2.5 * 1.1**1.5, outer_jvp) self.assertIsNone(acc.jvp(outer_acc.jvp(primal_out))) + @parameterized.product( + shared_first=[True, False], record_backprop=[True, False]) + def testFunctionJVPWithChangingTangentAliasing(self, shared_first, + record_backprop): + + @def_function.function + def multiply(x, y): + return x * y + + x = constant_op.constant(2.) + y = constant_op.constant(5.) + concrete = multiply.get_concrete_function(x, y) + cache_sizes = [] + for shared in (shared_first, not shared_first, shared_first): + tangent_x = constant_op.constant(2.) + tangent_y = tangent_x if shared else constant_op.constant(3.) + with backprop.GradientTape() as tape: + if record_backprop: + tape.watch((x, y)) + with forwardprop.ForwardAccumulator( + (x, y), (tangent_x, tangent_y)) as acc: + result = multiply(x, y) + jvp = acc.jvp(result) + self.assertAllClose(14. if shared else 16., jvp) + if record_backprop: + self.assertAllClose([5., 2.], tape.gradient(result, (x, y))) + cache_sizes.append( + len(concrete._first_order_tape_functions) + + len(concrete._higher_order_tape_functions)) + # Fresh tensors with the same aliasing pattern should reuse the cache. + self.assertEqual(cache_sizes[1], cache_sizes[2]) + + @parameterized.product( + dtype=[dtypes.float32, dtypes.float64], + transform=["identity", "sin", "square"], + compiled=[False, True]) + def testNestedForwardWithReusedIntermediate(self, dtype, transform, compiled): + transforms = { + "identity": array_ops.identity, + "sin": math_ops.sin, + "square": math_ops.square, + } + + def second_derivative(x): + with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as outer: + with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as inner: + m = transforms[transform](x) + result = (m * m) * m + first_derivative = inner.jvp(result) + return outer.jvp(first_derivative) + + if compiled: + second_derivative = def_function.function(second_derivative) + for value in (0.7, 1.1): + expected = { + "identity": 6 * value, + "sin": 6 * np.sin(value) * np.cos(value)**2 - 3 * np.sin(value)**3, + "square": 30 * value**4, + }[transform] + self.assertAllClose( + expected, second_derivative(constant_op.constant(value, dtype=dtype))) + + @parameterized.parameters(dtypes.float32, dtypes.float64) + def testNestedForwardWithReusedMatrixIntermediate(self, dtype): + x = constant_op.constant(0.7, dtype=dtype) + with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as outer: + with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as inner: + one = constant_op.constant(1., dtype=dtype) + m = array_ops_stack.stack([ + array_ops_stack.stack([x, one]), + array_ops_stack.stack([-one, x])]) + result = math_ops.matmul(math_ops.matmul(m, m), m) + first_derivative = inner.jvp(result) + self.assertAllClose(6 * m, outer.jvp(first_derivative)) + @test_util.assert_no_new_pyobjects_executing_eagerly() def testJVPPacking(self): two = constant_op.constant(2.) diff --git a/tensorflow/python/eager/polymorphic_function/concrete_function.py b/tensorflow/python/eager/polymorphic_function/concrete_function.py index a68acdd94d40e0..675663e4195e00 100644 --- a/tensorflow/python/eager/polymorphic_function/concrete_function.py +++ b/tensorflow/python/eager/polymorphic_function/concrete_function.py @@ -1590,9 +1590,19 @@ def _select_forward_and_backward_functions( input_tangents = forwardprop_util.TangentInfo() need_gradients_for_jvps = record.should_record_backprop( input_tangents.tangents) + forwardprop_input_aliases = () + if input_tangents.tangents: + # JVP wrappers capture inputs by tensor identity, so generated functions + # depend on which inputs and tangents alias. Normalize the IDs to keep + # cache entries reusable across calls with fresh tensors. + tensor_indices = {} + forwardprop_input_aliases = tuple( + tensor_indices.setdefault(ops.tensor_id(t), len(tensor_indices)) + for t in args + input_tangents.tangents) # Allows re-use of forward and backward function pairs depending on the # tapes and forward accumulators watching its inputs. - cache_key = (need_gradients_for_jvps, input_tangents.indices) + cache_key = (need_gradients_for_jvps, input_tangents.indices, + forwardprop_input_aliases) if (possible_gradient_type == gradients_util.POSSIBLE_GRADIENT_TYPES_FIRST_ORDER): if input_tangents.indices or executing_eagerly: From 409ead9f41a5789c921a574e9e9cdc421e66a1a5 Mon Sep 17 00:00:00 2001 From: Sai Ramineni Date: Fri, 25 Sep 2026 09:49:57 -0400 Subject: [PATCH 11/25] Guard GPU-only integer Pow header in CPU builds --- tensorflow/core/kernels/cwise_op_pow.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/core/kernels/cwise_op_pow.cc b/tensorflow/core/kernels/cwise_op_pow.cc index 73edeae5cfd520..40815b537d073d 100644 --- a/tensorflow/core/kernels/cwise_op_pow.cc +++ b/tensorflow/core/kernels/cwise_op_pow.cc @@ -14,7 +14,10 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/core/kernels/cwise_ops_common.h" + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM #include "tensorflow/core/kernels/integer_pow_gpu.h" +#endif namespace tensorflow { REGISTER6(BinaryOp, CPU, "Pow", functor::pow, float, Eigen::half, bfloat16, From fb411edf7f4f7947c294d4496c967b026b61930d Mon Sep 17 00:00:00 2001 From: Bhushan Asati Date: Fri, 25 Sep 2026 14:19:02 -0400 Subject: [PATCH 12/25] Iterate forward-mode cache inputs without concatenation --- .../python/eager/polymorphic_function/concrete_function.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/eager/polymorphic_function/concrete_function.py b/tensorflow/python/eager/polymorphic_function/concrete_function.py index 675663e4195e00..8f230526b4a1bc 100644 --- a/tensorflow/python/eager/polymorphic_function/concrete_function.py +++ b/tensorflow/python/eager/polymorphic_function/concrete_function.py @@ -16,6 +16,7 @@ """Implementation for ConcreteFunction.""" import collections +import itertools from tensorflow.core.framework import attr_value_pb2 from tensorflow.core.function.polymorphism import function_type as function_type_lib @@ -1598,7 +1599,7 @@ def _select_forward_and_backward_functions( tensor_indices = {} forwardprop_input_aliases = tuple( tensor_indices.setdefault(ops.tensor_id(t), len(tensor_indices)) - for t in args + input_tangents.tangents) + for t in itertools.chain(args, input_tangents.tangents)) # Allows re-use of forward and backward function pairs depending on the # tapes and forward accumulators watching its inputs. cache_key = (need_gradients_for_jvps, input_tangents.indices, From 6ad8816104a943e3a97904d20405b7d8066a6c1a Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Sun, 27 Sep 2026 04:54:16 -0700 Subject: [PATCH 13/25] [XLA] Upgrade nanobind to da42c8a5890eff09e672e7cfd953627655aa3e28 Fixes a data race seen in free threading builds. PiperOrigin-RevId: 989154130 --- third_party/xla/third_party/nanobind/workspace.bzl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/third_party/xla/third_party/nanobind/workspace.bzl b/third_party/xla/third_party/nanobind/workspace.bzl index 5a265d0b274a43..232b319104852c 100644 --- a/third_party/xla/third_party/nanobind/workspace.bzl +++ b/third_party/xla/third_party/nanobind/workspace.bzl @@ -20,8 +20,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): tf_http_archive( name = "nanobind", - strip_prefix = "nanobind-68480a9e6883bb1bf65e1925423321e2a89d07f9", - sha256 = "177ca257aacbf58004809f8cd1c20ad0416d8ace0a884c16a3c03ea87e81876b", - urls = tf_mirror_urls("https://github.com/wjakob/nanobind/archive/68480a9e6883bb1bf65e1925423321e2a89d07f9.tar.gz"), + strip_prefix = "nanobind-da42c8a5890eff09e672e7cfd953627655aa3e28", + sha256 = "3b24ad20ce0429b1a501151eeaeff723451b9e9b812887b7f9cb8575fa5e7ec7", + urls = tf_mirror_urls("https://github.com/wjakob/nanobind/archive/da42c8a5890eff09e672e7cfd953627655aa3e28.tar.gz"), build_file = "//third_party/nanobind:nanobind.BUILD", ) From 1a01b6463958c3420b7cb0739fa82cfe47abacfa Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Sun, 27 Sep 2026 08:03:52 -0700 Subject: [PATCH 14/25] [XLA:GPU][AMD] Disable topk_special_floats_test to unblock presubmits. Something is broken at head. PiperOrigin-RevId: 989195429 --- third_party/xla/xla/tests/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 184fac20a87fcc..f85b93f8da690d 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -2194,6 +2194,7 @@ xla_test( "cpu", "gpu", ], + disabled_backends = ["amdgpu_any"], deps = [ ":hlo_pjrt_interpreter_reference_mixin", ":hlo_test_base", From 659a0a51e6b55f4e68fc70a60da05404ac8a153f Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Sun, 27 Sep 2026 08:27:09 -0700 Subject: [PATCH 15/25] [XLA:CPU] Remove redundant passes from the new XTile compilation pipeline Clean up dead, redundant, and duplicate passes in `AddNewXtileToVectorPasses` and `AddNewVectorToLLVMPasses`: - Remove `RegisterOptimizationPasses` before `LowerXTileEntryPass` (it ran `LoopInvariantCodeMotion`, `SymbolDCE` before `xtile.entry_func` was lowered). - Remove `ConvertElementwise0DTensorToScalarPass`, because the VectorizeXtile.cc handles it anyway. - Move `StablehloTargetIndependentOptimizationPass` it was a no-op pass after the vectorization. - Remove duplicate `LowerToLLVMPass` invocation in `AddNewVectorToLLVMPasses`. PiperOrigin-RevId: 989201049 --- .../xla/xla/backends/cpu/codegen/fusion_compiler.cc | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc index 84507d32cedc7c..78b9682f3bd082 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc @@ -408,13 +408,15 @@ void AddXtileToVectorPasses(mlir::OpPassManager& pm, bool msan_enabled, void AddNewXtileToVectorPasses(mlir::OpPassManager& pm, int32_t vector_width) { pm.addPass(xtile::createVerifyLegalXTileOpsPass()); - emitters::RegisterOptimizationPasses(pm); - pm.addPass(xtile::createExpandXtileComplexOpsPass()); + pm.addNestedPass( + mlir::stablehlo::createStablehloTargetIndependentOptimizationPass()); + pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(xtile::createStablehloLowerToArithPass()); pm.addPass(xtile::createLegalizeUnsignedIntegersAsSignlessPass()); pm.addPass(cpu::createLegalizeNarrowFloatStoragePass()); pm.addPass(mlir::createCanonicalizerPass()); + pm.addPass(mlir::createCSEPass()); pm.addPass(cpu::createVectorizeXTilePass()); pm.addPass(cpu::createLowerXTileEntryPass( cpu::LowerXTileEntryPassOptions{/*prefer_vector_width=*/vector_width})); @@ -424,16 +426,10 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm, int32_t vector_width) { pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(mlir::createCSEPass()); - pm.addNestedPass( - mlir::stablehlo::createStablehloTargetIndependentOptimizationPass()); - - pm.addPass(mlir::createCanonicalizerPass()); pm.addNestedPass( mlir::vector::createLowerVectorMultiReductionPass( mlir::vector::VectorMultiReductionLowering::InnerParallel)); - pm.addPass(xtile::createConvertElementwise0DTensorToScalarPass()); - pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(mlir::createCSEPass()); } @@ -495,7 +491,6 @@ void AddNewVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, mlir::vector::VectorTransposeLowering::Shuffle16x16; pm.addPass(mlir::createConvertVectorToLLVMPass(options)); pm.addPass(cpu::createLowerMemRefBitcastPass()); - pm.addPass(cpu::createLowerToLLVMPass()); pm.addPass(mlir::memref::createExpandStridedMetadataPass()); pm.addPass(emitters::createSafeIntegerArithmeticPass()); From 5ad2263db8a16accf660dd755e88583da6cbf88e Mon Sep 17 00:00:00 2001 From: Alexandre Boulgakov Date: Sun, 27 Sep 2026 09:26:39 -0700 Subject: [PATCH 16/25] [XLA] Extend HLO kDot with block scaling and sparsity. PiperOrigin-RevId: 989215525 --- .../xla/xla/hlo/builder/xla_builder.cc | 41 ++++-- third_party/xla/xla/hlo/builder/xla_builder.h | 20 ++- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 32 ++++- third_party/xla/xla/hlo/ir/hlo_instruction.h | 7 + .../xla/xla/hlo/ir/hlo_instruction_test.cc | 39 ++++++ .../xla/xla/hlo/ir/hlo_instructions.cc | 45 +++++-- third_party/xla/xla/hlo/ir/hlo_instructions.h | 25 ++++ third_party/xla/xla/hlo/ir/hlo_opcode.h | 2 +- third_party/xla/xla/hlo/parser/hlo_parser.cc | 29 ++++- .../xla/xla/hlo/parser/hlo_parser_test.cc | 28 ++++ .../simplifiers/algebraic_simplifier.cc | 3 + .../simplifiers/batch_dot_simplification.cc | 3 + third_party/xla/xla/python/ops.cc | 14 +- .../xla/xla/service/hlo_creation_utils.cc | 32 ++++- .../xla/xla/service/hlo_creation_utils.h | 9 ++ third_party/xla/xla/service/hlo_verifier.cc | 123 +++++++++--------- .../xla/xla/service/hlo_verifier_test.cc | 59 ++++++++- .../xla/xla/service/shape_inference.cc | 64 ++++++--- third_party/xla/xla/service/shape_inference.h | 3 +- .../xla/xla/service/shape_inference_test.cc | 40 +++++- .../xla/xla/service/spmd/dot_handler.cc | 13 ++ .../xla/xla/service/transpose_folding.cc | 5 +- .../xla/xla/service/transpose_folding_test.cc | 10 +- 23 files changed, 510 insertions(+), 136 deletions(-) diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index fafc71bb7d0909..62434b4f056a9d 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -2191,15 +2191,20 @@ XlaOp XlaBuilder::Dot(XlaOp lhs, XlaOp rhs, XlaOp XlaBuilder::DotGeneral( XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig* precision_config, - std::optional preferred_element_type) { + std::optional preferred_element_type, + absl::Span ext_operands, const SparsityConfig* sparsity_config, + const BlockScalingConfig* block_scaling_config) { return ReportErrorOrReturn([&]() -> absl::StatusOr { ABSL_ASSIGN_OR_RETURN(const Shape* lhs_shape, GetShapePtr(lhs)); ABSL_ASSIGN_OR_RETURN(const Shape* rhs_shape, GetShapePtr(rhs)); - ABSL_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferDotOpShape( - *lhs_shape, *rhs_shape, dimension_numbers, - preferred_element_type)); + ABSL_ASSIGN_OR_RETURN( + Shape shape, + ShapeInference::InferDotOpShape( + *lhs_shape, *rhs_shape, dimension_numbers, preferred_element_type, + sparsity_config ? *sparsity_config : SparsityConfig())); return DotGeneralInternal(shape, lhs, rhs, dimension_numbers, - precision_config); + precision_config, ext_operands, sparsity_config, + block_scaling_config); }); } @@ -2229,14 +2234,26 @@ XlaOp XlaBuilder::ScaledDot( absl::StatusOr XlaBuilder::DotGeneralInternal( const Shape& shape, XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_numbers, - const PrecisionConfig* precision_config) { + const PrecisionConfig* precision_config, + absl::Span ext_operands, const SparsityConfig* sparsity_config, + const BlockScalingConfig* block_scaling_config) { HloInstructionProto instr; *instr.mutable_shape() = shape.ToProto(); *instr.mutable_dot_dimension_numbers() = dimension_numbers; if (precision_config != nullptr) { *instr.mutable_precision_config() = *precision_config; } - return AddInstruction(std::move(instr), HloOpcode::kDot, {lhs, rhs}); + if (sparsity_config && + (sparsity_config->has_lhs() || sparsity_config->has_rhs())) { + *instr.mutable_sparsity_config() = *sparsity_config; + } + if (block_scaling_config && + (block_scaling_config->has_lhs() || block_scaling_config->has_rhs())) { + *instr.mutable_block_scaling_config() = *block_scaling_config; + } + std::vector operands = {lhs, rhs}; + operands.insert(operands.end(), ext_operands.begin(), ext_operands.end()); + return AddInstruction(std::move(instr), HloOpcode::kDot, operands); } XlaOp ScaledDot(const XlaOp lhs, const XlaOp rhs, const XlaOp lhs_scale, @@ -5814,9 +5831,13 @@ XlaOp Dot(const XlaOp lhs, const XlaOp rhs, XlaOp DotGeneral(const XlaOp lhs, const XlaOp rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig* precision_config, - std::optional preferred_element_type) { - return lhs.builder()->DotGeneral(lhs, rhs, dimension_numbers, - precision_config, preferred_element_type); + std::optional preferred_element_type, + absl::Span ext_operands, + const SparsityConfig* sparsity_config, + const BlockScalingConfig* block_scaling_config) { + return lhs.builder()->DotGeneral( + lhs, rhs, dimension_numbers, precision_config, preferred_element_type, + ext_operands, sparsity_config, block_scaling_config); } XlaOp RaggedAllToAll(const XlaOp input, const XlaOp input_offsets, diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index 7cceabb9adc353..980fa0a78fcebd 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -709,7 +709,10 @@ class XlaBuilder { XlaOp DotGeneral( XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig* precision_config = nullptr, - std::optional preferred_element_type = std::nullopt); + std::optional preferred_element_type = std::nullopt, + absl::Span ext_operands = {}, + const SparsityConfig* sparsity_config = nullptr, + const BlockScalingConfig* block_scaling_config = nullptr); XlaOp RaggedAllToAll( XlaOp input, XlaOp input_offsets, XlaOp send_sizes, XlaOp output, @@ -1534,7 +1537,10 @@ class XlaBuilder { friend XlaOp DotGeneral(XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_number, const PrecisionConfig* precision_config, - std::optional preferred_element_type); + std::optional preferred_element_type, + absl::Span ext_operands, + const SparsityConfig* sparsity_config, + const BlockScalingConfig* block_scaling_config); friend XlaOp RaggedDot(XlaOp lhs, XlaOp rhs, XlaOp group_sizes, const RaggedDotDimensionNumbers& dimension_numbers, const PrecisionConfig* precision_config, @@ -1546,7 +1552,10 @@ class XlaBuilder { virtual absl::StatusOr DotGeneralInternal( const Shape& shape, XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_number, - const PrecisionConfig* precision_config); + const PrecisionConfig* precision_config, + absl::Span ext_operands, + const SparsityConfig* sparsity_config, + const BlockScalingConfig* block_scaling_config); friend XlaOp RaggedAllToAll(XlaOp input, XlaOp input_offsets, XlaOp send_sizes, XlaOp output, XlaOp output_offsets, XlaOp recv_sizes, @@ -2598,7 +2607,10 @@ XlaOp Dot(XlaOp lhs, XlaOp rhs, XlaOp DotGeneral( XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig* precision_config = nullptr, - std::optional preferred_element_type = std::nullopt); + std::optional preferred_element_type = std::nullopt, + absl::Span ext_operands = {}, + const SparsityConfig* sparsity_config = nullptr, + const BlockScalingConfig* block_scaling_config = nullptr); // Enqueues a ragged all to all instruction onto the computation. XlaOp RaggedAllToAll( diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 3563727f001b7a..e0644086582f5c 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -1235,10 +1235,10 @@ absl::StatusOr> HloInstruction::CreateFromProto( PrecisionConfig precision_config = proto.precision_config(); precision_config.mutable_operand_precision()->Resize( proto.operand_ids_size(), PrecisionConfig::DEFAULT); - auto operand_vector = all_operands(); instruction = std::make_unique( - shape, operands(0), operands(1), proto.dot_dimension_numbers(), - precision_config); + shape, all_operands(), proto.dot_dimension_numbers(), + precision_config, proto.sparsity_config(), + proto.block_scaling_config()); break; } case HloOpcode::kRaggedDot: { @@ -1797,8 +1797,18 @@ HloInstruction::CreateTriangularSolve(const Shape& shape, HloInstruction* a, const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config) { - return std::make_unique(shape, lhs, rhs, dimension_numbers, - precision_config); + return CreateDot(shape, {lhs, rhs}, dimension_numbers, precision_config); +} + +/* static */ std::unique_ptr HloInstruction::CreateDot( + const Shape& shape, absl::Span operands, + const DotDimensionNumbers& dimension_numbers, + const PrecisionConfig& precision_config, + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config) { + return std::make_unique(shape, operands, dimension_numbers, + precision_config, sparsity_config, + block_scaling_config); } /* static */ std::unique_ptr HloInstruction::CreateRaggedDot( @@ -6459,20 +6469,32 @@ const RaggedDotDimensionNumbers& HloInstruction::ragged_dot_dimension_numbers() } const SparsityConfig& HloInstruction::sparsity_config() const { + if (auto d = DynCast(this)) { + return d->sparsity_config(); + } return Cast(this)->sparsity_config(); } void HloInstruction::set_sparsity_config( const SparsityConfig& sparsity_config) { + if (auto d = DynCast(this)) { + return d->set_sparsity_config(sparsity_config); + } Cast(this)->set_sparsity_config(sparsity_config); } const BlockScalingConfig& HloInstruction::block_scaling_config() const { + if (auto d = DynCast(this)) { + return d->block_scaling_config(); + } return Cast(this)->block_scaling_config(); } void HloInstruction::set_block_scaling_config( const BlockScalingConfig& block_scaling_config) { + if (auto d = DynCast(this)) { + return d->set_block_scaling_config(block_scaling_config); + } Cast(this)->set_block_scaling_config( block_scaling_config); } diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index c233a089e30a74..0b6ff3b7fa83e7 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -494,6 +494,13 @@ class HloInstruction { const DotDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config); + static std::unique_ptr CreateDot( + const Shape& shape, absl::Span operands, + const DotDimensionNumbers& dimension_numbers, + const PrecisionConfig& precision_config, + const SparsityConfig& sparsity_config = {}, + const BlockScalingConfig& block_scaling_config = {}); + // Creates a ragged dot op with operands 'lhs', 'rhs', and 'group_sizes', with // contracting, batch, ragged, and group dimensions specified in // 'dimension_numbers'. diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc index 6ed7ab9dc268b0..07cdbdc9ee3ceb 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc @@ -207,6 +207,45 @@ TEST_F(HloInstructionTest, BlockScalingConfigToString) { } } +TEST_F(HloInstructionTest, DotBlockScalingAndSparsityToString) { + HloComputation::Builder builder("main"); + auto lhs = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(BF16, {64, 64}), "lhs")); + auto rhs = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(BF16, {128, 64}), "rhs")); + auto lhs_scale = builder.AddInstruction(HloInstruction::CreateParameter( + 2, ShapeUtil::MakeShape(F8E8M0FNU, {64, 2}), "lhs_scale")); + auto lhs_indices = builder.AddInstruction(HloInstruction::CreateParameter( + 3, ShapeUtil::MakeShape(S8, {64, 16}), "lhs_indices")); + + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(1); + dnums.add_rhs_contracting_dimensions(0); + PrecisionConfig precision_config; + + SparsityConfig sp; + sp.mutable_lhs()->set_idx(3); + sp.mutable_lhs()->set_num_non_zero(2); + sp.mutable_lhs()->set_block_size(4); + sp.mutable_lhs()->set_dimension(1); + sp.mutable_lhs()->set_stride(1); + + BlockScalingConfig bs; + bs.mutable_lhs()->set_scale_idx(2); + bs.mutable_lhs()->add_strides(1); + bs.mutable_lhs()->add_strides(32); + bs.mutable_lhs()->add_steps(1); + bs.mutable_lhs()->add_steps(1); + + auto dot = builder.AddInstruction(HloInstruction::CreateDot( + ShapeUtil::MakeShape(BF16, {64, 64}), {lhs, rhs, lhs_scale, lhs_indices}, + dnums, precision_config, sp, bs)); + + EXPECT_EQ( + dot->ToString(), + R"(%dot = bf16[64,64]{1,0} dot(%lhs, %rhs, %lhs_scale, %lhs_indices), lhs_contracting_dims={1}, rhs_contracting_dims={0}, sparsity_config={lhs={sparsity=2x4 dimension=1 stride=1 idx=3}}, block_scaling_config={lhs={scale_idx=2 strides=1x32 steps=1x1}})"); +} + TEST_F(HloInstructionTest, GetStackTraceStringFromStackFrameId) { auto module = CreateNewVerifiedModule(); HloComputation::Builder builder("main"); diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index 4071f0c416a838..66516ba477ae27 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -4546,17 +4546,30 @@ HloDotInstruction::HloDotInstruction( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config) + : HloDotInstruction(shape, {lhs, rhs}, dimension_numbers, precision_config, + SparsityConfig(), BlockScalingConfig()) {} + +HloDotInstruction::HloDotInstruction( + const Shape& shape, absl::Span operands, + const DotDimensionNumbers& dimension_numbers, + const PrecisionConfig& precision_config, + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config) : HloInstruction(HloOpcode::kDot, shape), dot_dimension_numbers_(dimension_numbers), - precision_config_(precision_config) { - AppendOperand(lhs); - AppendOperand(rhs); + precision_config_(precision_config), + sparsity_config_(sparsity_config), + block_scaling_config_(block_scaling_config) { + CHECK_GE(operands.size(), 2); + AppendOperands(operands); } void HloDotInstruction::ToProto(HloInstructionProto* proto) const { HloInstruction::ToProto(proto); *proto->mutable_dot_dimension_numbers() = dot_dimension_numbers_; *proto->mutable_precision_config() = precision_config_; + *proto->mutable_sparsity_config() = sparsity_config_; + *proto->mutable_block_scaling_config() = block_scaling_config_; } void HloDotInstruction::PrintExtraAttributesImpl( @@ -4565,6 +4578,18 @@ void HloDotInstruction::PrintExtraAttributesImpl( printer->Append(DotDimensionNumbersToString(dot_dimension_numbers_)); }); PrintPrecisionConfig(printer, precision_config_); + auto print_config = [&printer](absl::string_view name, + const std::string& config) { + printer.Next( + [&](Printer* p) { p->Append(absl::StrCat(name, "={", config, "}")); }); + }; + if (sparsity_config_.has_lhs() || sparsity_config_.has_rhs()) { + print_config("sparsity_config", SparsityConfigToString(sparsity_config_)); + } + if (block_scaling_config_.has_lhs() || block_scaling_config_.has_rhs()) { + print_config("block_scaling_config", + BlockScalingConfigToString(block_scaling_config_)); + } } bool HloDotInstruction::IdenticalSlowPath( @@ -4574,17 +4599,21 @@ bool HloDotInstruction::IdenticalSlowPath( const auto& casted_other = static_cast(other); return protobuf_util::HaveSameSerialization( dot_dimension_numbers(), casted_other.dot_dimension_numbers()) && - protobuf_util::HaveSameSerialization(precision_config(), - casted_other.precision_config()); + protobuf_util::HaveSameSerialization( + precision_config(), casted_other.precision_config()) && + protobuf_util::HaveSameSerialization(sparsity_config(), + casted_other.sparsity_config()) && + protobuf_util::HaveSameSerialization( + block_scaling_config(), casted_other.block_scaling_config()); } std::unique_ptr HloDotInstruction::CloneWithNewOperandsImpl( const Shape& shape, absl::Span new_operands, HloCloneContext* context) const { - CHECK_EQ(new_operands.size(), 2); + CHECK_GE(new_operands.size(), 2); return std::make_unique( - shape, new_operands[0], new_operands[1], dot_dimension_numbers_, - precision_config_); + shape, new_operands, dot_dimension_numbers_, precision_config_, + sparsity_config_, block_scaling_config_); } HloRaggedDotInstruction::HloRaggedDotInstruction( diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index f67a1714749161..d1e46162ccd829 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -2816,6 +2816,25 @@ class HloDotInstruction : public HloInstruction { const DotDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config); + explicit HloDotInstruction(const Shape& shape, + absl::Span operands, + const DotDimensionNumbers& dimension_numbers, + const PrecisionConfig& precision_config, + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config); + + const SparsityConfig& sparsity_config() const { return sparsity_config_; } + void set_sparsity_config(const SparsityConfig& sparsity_config) { + sparsity_config_ = sparsity_config; + } + + const BlockScalingConfig& block_scaling_config() const { + return block_scaling_config_; + } + void set_block_scaling_config(const BlockScalingConfig& config) { + block_scaling_config_ = config; + } + // Returns data on the dimension numbers used for a dot operation. const DotDimensionNumbers& dot_dimension_numbers() const { return dot_dimension_numbers_; @@ -2860,6 +2879,12 @@ class HloDotInstruction : public HloInstruction { // Information used to communicate to the implementation about the algorithm // used to produce results. See the documentation on precision_config(). PrecisionConfig precision_config_; + + // The sparsity configuration used for the dot. + SparsityConfig sparsity_config_; + + // Dot block scaling config. + BlockScalingConfig block_scaling_config_; }; class HloRaggedDotInstruction : public HloInstruction { diff --git a/third_party/xla/xla/hlo/ir/hlo_opcode.h b/third_party/xla/xla/hlo/ir/hlo_opcode.h index 9da5dc0342a2d7..cb8f668150bcff 100644 --- a/third_party/xla/xla/hlo/ir/hlo_opcode.h +++ b/third_party/xla/xla/hlo/ir/hlo_opcode.h @@ -100,7 +100,7 @@ namespace xla { V(kCustomCall, "custom-call", kHloOpcodeIsVariadic) \ V(kDivide, "divide", 2) \ V(kDomain, "domain", 1) \ - V(kDot, "dot", 2) \ + V(kDot, "dot", kHloOpcodeIsVariadic) \ V(kDynamicReshape, "dynamic-reshape", kHloOpcodeIsVariadic) \ V(kDynamicSlice, "dynamic-slice", kHloOpcodeIsVariadic) \ V(kDynamicUpdateSlice, "dynamic-update-slice", kHloOpcodeIsVariadic) \ diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index b6e173aeb71a64..37d0071dbbd980 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -3724,10 +3724,22 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT attrs["algorithm"] = {/*required=*/false, AttrTy::kPrecisionAlgorithm, &algorithm}; + optional parsed_sparsity_config; + attrs["sparsity_config"] = {/*required=*/false, AttrTy::kSparsityConfig, + &parsed_sparsity_config}; + optional parsed_block_scaling_config; + attrs["block_scaling_config"] = {/*required=*/false, + AttrTy::kBlockScalingConfig, + &parsed_block_scaling_config}; + if ((!preset_operands && !ParseOperands(&operands, builder)) || !ParseAttributes(attrs, allow_attributes, shape)) { return nullptr; } + if (operands.size() < 2) { + Error(lexer_.GetLoc(), "expects at least 2 operands"); + return nullptr; + } DotDimensionNumbers dnum; if (lhs_contracting_dims) { @@ -3758,15 +3770,17 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT if (algorithm) { precision_config.set_algorithm(*algorithm); } + SparsityConfig sp = parsed_sparsity_config.value_or(SparsityConfig()); if (!maybe_infer_shape([&] { return ShapeInference::InferDotOpShape( operands[0]->shape(), operands[1]->shape(), dnum, - /*preferred_element_type=*/std::nullopt); + /*preferred_element_type=*/std::nullopt, sp); })) { return nullptr; } return builder->AddInstruction(HloInstruction::CreateDot( - *shape, operands[0], operands[1], dnum, precision_config)); + *shape, operands, dnum, precision_config, sp, + parsed_block_scaling_config.value_or(BlockScalingConfig()))); } case HloOpcode::kRaggedDot: { optional> lhs_contracting_dims; @@ -9240,10 +9254,13 @@ HloComputation* HloParserImpl::CreateAsyncWrappedComputation( // will fail and crash. When there are not enough operands at creation time // (when late binding is used), we add dummy operands, and update the // async-wrapped computation later. - std::optional async_wrapped_opcode_arity = - HloOpcodeArity(async_wrapped_opcode); - uint64_t num_async_operands = std::max( - async_wrapped_opcode_arity.value_or(0), operand_shapes.size()); + int64_t min_arity = HloOpcodeArity(async_wrapped_opcode).value_or(0); + if (async_wrapped_opcode == HloOpcode::kDot || + async_wrapped_opcode == HloOpcode::kConvolution) { + min_arity = 2; + } + uint64_t num_async_operands = + std::max(min_arity, operand_shapes.size()); HloComputation::Builder async_wrapped_builder("async_wrapped"); async_wrapped_operands.reserve(num_async_operands); diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc index ddf4306437a436..6578a8cc02e1e4 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -7494,6 +7494,34 @@ ENTRY BlockScalingConfig { EXPECT_EQ(config_after.DebugString(), config_before.DebugString()); } +TEST_F(HloParserTest, DotBlockScalingAndSparsityConfig_RoundTrip) { + const char* const hlo_string = R"( +HloModule DotBlockScalingConfigModule +ENTRY DotBlockScalingConfig { + %lhs = bf16[64,64] parameter(0) + %rhs = bf16[128,64] parameter(1) + %lhs_scale = f8e8m0fnu[64,2] parameter(2) + %lhs_indices = s8[64,16] parameter(3) + ROOT %dot = bf16[64,64] dot(%lhs, %rhs, %lhs_scale, %lhs_indices), + lhs_contracting_dims={1}, rhs_contracting_dims={0}, + sparsity_config={lhs={sparsity=2x4 dimension=1 stride=1 idx=3}}, + block_scaling_config={lhs={scale_idx=2 strides=1x32 steps=1x1}} +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); + auto* dot_before = module->entry_computation()->root_instruction(); + auto bs_before = dot_before->block_scaling_config(); + auto sp_before = dot_before->sparsity_config(); + std::string printed = module->ToString(); + ASSERT_OK_AND_ASSIGN(auto parsed_module, + ParseAndReturnUnverifiedModule(printed)); + auto* dot_after = parsed_module->entry_computation()->root_instruction(); + EXPECT_EQ(dot_after->block_scaling_config().DebugString(), + bs_before.DebugString()); + EXPECT_EQ(dot_after->sparsity_config().DebugString(), + sp_before.DebugString()); +} + TEST_F(HloParserTest, DesugarParsingTest_DotStart) { const char* const hlo = R"( HloModule async_dot_example diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 2193ddb2c91c6e..da7187f072fbca 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -4201,6 +4201,9 @@ AlgebraicSimplifierVisitor::MakeMultiplyForPrecisionAlgorithm( absl::Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { CHECK(computation_ == dot->parent()); + if (dot->operand_count() > 2) { + return absl::OkStatus(); + } HloDotInstruction* dot_cast = Cast(dot); const auto& dnums = dot->dot_dimension_numbers(); diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/batch_dot_simplification.cc b/third_party/xla/xla/hlo/transforms/simplifiers/batch_dot_simplification.cc index c9df182aa9d79f..b5c18cfc7123f0 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/batch_dot_simplification.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/batch_dot_simplification.cc @@ -41,6 +41,9 @@ namespace xla { absl::StatusOr BatchDotSimplification::ElideDegenerateBatchDimensionFromBatchDot( HloInstruction* batch_dot) { + if (batch_dot->operand_count() > 2) { + return false; + } // This pass assumes the lhs and rhs batch dimensions are equal and strictly // ascending. const auto& is_iota = [](absl::Span dims) { diff --git a/third_party/xla/xla/python/ops.cc b/third_party/xla/xla/python/ops.cc index 1c50da9ac2127d..0bf1b0f63c854c 100644 --- a/third_party/xla/xla/python/ops.cc +++ b/third_party/xla/xla/python/ops.cc @@ -570,9 +570,17 @@ NB_MODULE(_ops, m) { m.def("Dot", &Dot, nb::arg("lhs"), nb::arg("rhs"), nb::arg("precision_config") = nullptr, nb::arg("preferred_element_type") = std::nullopt); - m.def("DotGeneral", &DotGeneral, nb::arg("lhs"), nb::arg("rhs"), - nb::arg("dimension_numbers"), nb::arg("precision_config") = nullptr, - nb::arg("preferred_element_type") = std::nullopt); + m.def( + "DotGeneral", + [](XlaOp lhs, XlaOp rhs, const DotDimensionNumbers& dimension_numbers, + const PrecisionConfig* precision_config, + std::optional preferred_element_type) { + return DotGeneral(lhs, rhs, dimension_numbers, precision_config, + preferred_element_type); + }, + nb::arg("lhs"), nb::arg("rhs"), nb::arg("dimension_numbers"), + nb::arg("precision_config") = nullptr, + nb::arg("preferred_element_type") = std::nullopt); m.def("DynamicReshape", static_cast, absl::Span, diff --git a/third_party/xla/xla/service/hlo_creation_utils.cc b/third_party/xla/xla/service/hlo_creation_utils.cc index 0b5dcca89d7c33..38a0e6cd699d2d 100644 --- a/third_party/xla/xla/service/hlo_creation_utils.cc +++ b/third_party/xla/xla/service/hlo_creation_utils.cc @@ -379,14 +379,32 @@ absl::StatusOr MakeDotHlo( const PrecisionConfig& precision_config, std::optional preferred_element_type, const OpMetadata* metadata) { - HloComputation* computation = lhs->parent(); - CHECK_EQ(computation, rhs->parent()); - ABSL_ASSIGN_OR_RETURN(Shape dot_shape, ShapeInference::InferDotOpShape( - lhs->shape(), rhs->shape(), dim_numbers, - preferred_element_type)); + return MakeDotHlo({lhs, rhs}, dim_numbers, precision_config, + preferred_element_type, SparsityConfig(), + BlockScalingConfig(), metadata); +} + +absl::StatusOr MakeDotHlo( + absl::Span operands, + const DotDimensionNumbers& dim_numbers, + const PrecisionConfig& precision_config, + std::optional preferred_element_type, + const SparsityConfig& sparsity_config, + const BlockScalingConfig& block_scaling_config, + const OpMetadata* metadata) { + CHECK_GE(operands.size(), 2); + HloComputation* computation = operands[0]->parent(); + for (HloInstruction* operand : operands) { + CHECK_EQ(computation, operand->parent()); + } + ABSL_ASSIGN_OR_RETURN(Shape dot_shape, + ShapeInference::InferDotOpShape( + operands[0]->shape(), operands[1]->shape(), dim_numbers, + preferred_element_type, sparsity_config)); return computation->AddInstruction( - HloInstruction::CreateDot(dot_shape, lhs, rhs, dim_numbers, - precision_config), + HloInstruction::CreateDot(dot_shape, operands, dim_numbers, + precision_config, sparsity_config, + block_scaling_config), metadata); } diff --git a/third_party/xla/xla/service/hlo_creation_utils.h b/third_party/xla/xla/service/hlo_creation_utils.h index ea4805ec25deb7..e534d7833482ef 100644 --- a/third_party/xla/xla/service/hlo_creation_utils.h +++ b/third_party/xla/xla/service/hlo_creation_utils.h @@ -184,6 +184,15 @@ absl::StatusOr MakeDotHlo( std::optional preferred_element_type, const OpMetadata* metadata = nullptr); +absl::StatusOr MakeDotHlo( + absl::Span operands, + const DotDimensionNumbers& dim_numbers, + const PrecisionConfig& precision_config, + std::optional preferred_element_type, + const SparsityConfig& sparsity_config = {}, + const BlockScalingConfig& block_scaling_config = {}, + const OpMetadata* metadata = nullptr); + // Creates a RaggedDot HLO instruction and adds it to the computation containing // `lhs`, `rhs`, and `group_sizes` (all must be in the same computation). An // optional preferred_element_type can be specified to override the element diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index 0e532bfea16442..5ad67d8a5ae31e 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -191,12 +191,73 @@ absl::Status ShapeVerifier::HandleCopy(HloInstruction* copy) { return CheckUnaryShape(copy); } +absl::Status VerifySparsityAndBlockScaling(const HloInstruction* hlo) { + if (hlo->operand_count() < 2) { + return InvalidArgument("%s must have at least 2 operands, got %d", + HloOpcodeString(hlo->opcode()), + hlo->operand_count()); + } + std::vector seen_indices(hlo->operand_count(), false); + int64_t seen_count = 0; + auto check_idx = [&](int32_t idx, absl::string_view desc) -> absl::Status { + if (idx < 2 || idx >= hlo->operand_count()) { + return InvalidArgument("%s %d out of bounds", desc, idx); + } + if (!hlo->operand(idx)->shape().IsArray()) { + return InvalidArgument( + "Expected array argument for %s at index %d, but got %s", desc, idx, + ShapeUtil::HumanString(hlo->operand(idx)->shape())); + } + if (seen_indices[idx]) { + return InvalidArgument("Duplicate index %d for %s", idx, desc); + } + seen_indices[idx] = true; + ++seen_count; + return absl::OkStatus(); + }; + + if (hlo->sparsity_config().has_lhs()) { + ABSL_RETURN_IF_ERROR( + check_idx(hlo->sparsity_config().lhs().idx(), "Sparsity idx for lhs")); + } + if (hlo->sparsity_config().has_rhs()) { + ABSL_RETURN_IF_ERROR( + check_idx(hlo->sparsity_config().rhs().idx(), "Sparsity idx for rhs")); + } + if (hlo->block_scaling_config().has_lhs()) { + ABSL_RETURN_IF_ERROR(check_idx(hlo->block_scaling_config().lhs().scale_idx(), + "Block scaling scale_idx for lhs")); + if (hlo->block_scaling_config().lhs().has_zero_idx()) { + ABSL_RETURN_IF_ERROR(check_idx(hlo->block_scaling_config().lhs().zero_idx(), + "Block scaling zero_idx for lhs")); + } + } + if (hlo->block_scaling_config().has_rhs()) { + ABSL_RETURN_IF_ERROR(check_idx(hlo->block_scaling_config().rhs().scale_idx(), + "Block scaling scale_idx for rhs")); + if (hlo->block_scaling_config().rhs().has_zero_idx()) { + ABSL_RETURN_IF_ERROR(check_idx(hlo->block_scaling_config().rhs().zero_idx(), + "Block scaling zero_idx for rhs")); + } + } + if (seen_count != hlo->operand_count() - 2) { + return InvalidArgument( + "Expected all %d extra operands to be referenced by sparsity_config or " + "block_scaling_config, but %d were referenced", + hlo->operand_count() - 2, seen_count); + } + + return absl::OkStatus(); +} + absl::Status ShapeVerifier::HandleDot(HloInstruction* dot) { + ABSL_RETURN_IF_ERROR(VerifySparsityAndBlockScaling(dot)); ABSL_ASSIGN_OR_RETURN(const Shape expected, ShapeInference::InferDotOpShape( dot->operand(0)->shape(), dot->operand(1)->shape(), dot->dot_dimension_numbers(), - /*preferred_element_type=*/dot->shape().element_type())); + /*preferred_element_type=*/dot->shape().element_type(), + dot->sparsity_config())); return CheckShape(dot, expected); } @@ -303,65 +364,7 @@ absl::Status ShapeVerifier::HandleScaledDot(HloInstruction* scaled_dot) { } absl::Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) { - auto check_idx_in_range = [&](int32_t idx, int32_t low, int32_t high, - absl::string_view desc) -> absl::Status { - if (idx < low || idx >= high) { - return InvalidArgument("%s %d out of bounds", desc, idx); - } - return absl::OkStatus(); - }; - - if (convolution->sparsity_config().has_lhs()) { - ABSL_RETURN_IF_ERROR(check_idx_in_range( - convolution->sparsity_config().lhs().idx(), 2, - convolution->operand_count(), "Sparsity idx for lhs")); - } - if (convolution->sparsity_config().has_rhs()) { - ABSL_RETURN_IF_ERROR(check_idx_in_range( - convolution->sparsity_config().rhs().idx(), 2, - convolution->operand_count(), "Sparsity idx for rhs")); - } - if (convolution->block_scaling_config().has_lhs()) { - ABSL_RETURN_IF_ERROR(check_idx_in_range( - convolution->block_scaling_config().lhs().scale_idx(), 2, - convolution->operand_count(), "Block scaling scale_idx for lhs")); - if (convolution->block_scaling_config().lhs().has_zero_idx()) { - ABSL_RETURN_IF_ERROR(check_idx_in_range( - convolution->block_scaling_config().lhs().zero_idx(), 2, - convolution->operand_count(), "Block scaling zero_idx for lhs")); - if (convolution->block_scaling_config().lhs().scale_idx() == - convolution->block_scaling_config().lhs().zero_idx()) { - return InvalidArgument( - "LHS block scaling scale_idx and zero_idx cannot be the same (%d)", - convolution->block_scaling_config().lhs().scale_idx()); - } - } - } - if (convolution->block_scaling_config().has_rhs()) { - ABSL_RETURN_IF_ERROR(check_idx_in_range( - convolution->block_scaling_config().rhs().scale_idx(), 2, - convolution->operand_count(), "Block scaling scale_idx for rhs")); - if (convolution->block_scaling_config().rhs().has_zero_idx()) { - ABSL_RETURN_IF_ERROR(check_idx_in_range( - convolution->block_scaling_config().rhs().zero_idx(), 2, - convolution->operand_count(), "Block scaling zero_idx for rhs")); - if (convolution->block_scaling_config().rhs().scale_idx() == - convolution->block_scaling_config().rhs().zero_idx()) { - return InvalidArgument( - "RHS block scaling scale_idx and zero_idx cannot be the same (%d)", - convolution->block_scaling_config().rhs().scale_idx()); - } - } - } - if (convolution->block_scaling_config().has_lhs() && - convolution->block_scaling_config().has_rhs()) { - if (convolution->block_scaling_config().lhs().scale_idx() == - convolution->block_scaling_config().rhs().scale_idx()) { - return InvalidArgument( - "LHS and RHS block scaling scale_idx cannot be the same (%d)", - convolution->block_scaling_config().lhs().scale_idx()); - } - } + ABSL_RETURN_IF_ERROR(VerifySparsityAndBlockScaling(convolution)); ABSL_ASSIGN_OR_RETURN( Shape expected, diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc index b459ed7ad29e8f..e8cc45891e24ce 100644 --- a/third_party/xla/xla/service/hlo_verifier_test.cc +++ b/third_party/xla/xla/service/hlo_verifier_test.cc @@ -990,10 +990,8 @@ TEST_F(HloVerifierTest, ConvBlockScalingConfigSameScaleAndZeroIdxNotAllowed) { config.mutable_lhs()->set_zero_idx(2); conv->set_block_scaling_config(config); - EXPECT_THAT( - verifier().Run(module.get()).status().message(), - HasSubstr( - "LHS block scaling scale_idx and zero_idx cannot be the same (2)")); + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Duplicate index 2")); } TEST_F(HloVerifierTest, ConvBlockScalingConfigLhsAndRhsSameScaleIdxNotAllowed) { @@ -1005,9 +1003,56 @@ TEST_F(HloVerifierTest, ConvBlockScalingConfigLhsAndRhsSameScaleIdxNotAllowed) { config.mutable_rhs()->set_scale_idx(2); conv->set_block_scaling_config(config); - EXPECT_THAT( - verifier().Run(module.get()).status().message(), - HasSubstr("LHS and RHS block scaling scale_idx cannot be the same (2)")); + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Duplicate index 2")); +} + +static const char* const kDotWith4OperandsHloString = R"( +HloModule module +ENTRY entry_computation { + param0 = bf16[64,128] parameter(0) + param1 = bf16[128,64] parameter(1) + param2 = f8e8m0fnu[64,4] parameter(2) + param3 = f8e8m0fnu[4,64] parameter(3) + ROOT dot = bf16[64,64] dot(param0, param1, param2, param3), + lhs_contracting_dims={1}, rhs_contracting_dims={0} +})"; + +TEST_F(HloVerifierTest, DotBlockScalingConfigScaleIdxOutOfBoundsNotAllowed) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnUnverifiedModule(kDotWith4OperandsHloString)); + auto* dot = module->entry_computation()->root_instruction(); + BlockScalingConfig config; + config.mutable_rhs()->set_scale_idx(5); + dot->set_block_scaling_config(config); + + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Block scaling scale_idx for rhs 5 out of bounds")); +} + +TEST_F(HloVerifierTest, DotSparsityConfigLhsAndRhsSameIdxNotAllowed) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnUnverifiedModule(kDotWith4OperandsHloString)); + auto* dot = module->entry_computation()->root_instruction(); + SparsityConfig config; + config.mutable_lhs()->set_idx(2); + config.mutable_rhs()->set_idx(2); + dot->set_sparsity_config(config); + + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Duplicate index 2")); +} + +TEST_F(HloVerifierTest, DotUnreferencedExtraOperandsNotAllowed) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnUnverifiedModule(kDotWith4OperandsHloString)); + auto* dot = module->entry_computation()->root_instruction(); + BlockScalingConfig config; + config.mutable_lhs()->set_scale_idx(2); + dot->set_block_scaling_config(config); + + EXPECT_THAT(verifier().Run(module.get()).status().message(), + HasSubstr("Expected all 2 extra operands to be referenced")); } static const char* const kAddWithLayoutChangeHlo = R"( diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc index ed0006136a6cd3..d66e89f6d57a7e 100644 --- a/third_party/xla/xla/service/shape_inference.cc +++ b/third_party/xla/xla/service/shape_inference.cc @@ -857,9 +857,27 @@ absl::Status ValidateDotDimensionNumbers( return absl::OkStatus(); } +absl::StatusOr GetSparseDimSize( + int64_t raw_size, const SparsityConfig::TensorSparsityConfig& sp) { + if (sp.num_non_zero() <= 0 || sp.block_size() <= 0) { + return InvalidArgument( + "num_non_zero and block_size must be positive in SparsityConfig"); + } + if (IsUnboundedDynamicSize(raw_size)) { + return raw_size; + } + if (raw_size % sp.num_non_zero() != 0) { + return InvalidArgument( + "Sparse dimension size (%d) must be divisible by num_non_zero (%d)", + raw_size, sp.num_non_zero()); + } + return raw_size * sp.block_size() / sp.num_non_zero(); +} + absl::Status CheckDotDimensionConstraints( const Shape& lhs, const Shape& rhs, - const DotDimensionNumbers& dimension_numbers) { + const DotDimensionNumbers& dimension_numbers, + const SparsityConfig& sparsity_config = {}) { auto fail = [lhs, rhs](const std::string& addendum) -> absl::Status { std::string message = StrFormat("Cannot infer shape for dot operation: %s %s.", @@ -884,8 +902,19 @@ absl::Status CheckDotDimensionConstraints( dimension_numbers.lhs_contracting_dimensions(i); const int64_t rhs_contracting_dimension = dimension_numbers.rhs_contracting_dimensions(i); - if (!CompatibleDimensionSizes(lhs.dimensions(lhs_contracting_dimension), - rhs.dimensions(rhs_contracting_dimension))) { + int64_t lhs_size = lhs.dimensions(lhs_contracting_dimension); + int64_t rhs_size = rhs.dimensions(rhs_contracting_dimension); + if (sparsity_config.has_lhs() && + lhs_contracting_dimension == sparsity_config.lhs().dimension()) { + ABSL_ASSIGN_OR_RETURN(lhs_size, + GetSparseDimSize(lhs_size, sparsity_config.lhs())); + } + if (sparsity_config.has_rhs() && + rhs_contracting_dimension == sparsity_config.rhs().dimension()) { + ABSL_ASSIGN_OR_RETURN(rhs_size, + GetSparseDimSize(rhs_size, sparsity_config.rhs())); + } + if (!CompatibleDimensionSizes(lhs_size, rhs_size)) { return fail("Contracting dimension sizes are not compatible."); } } @@ -954,7 +983,8 @@ void GenerateDotResultDimensions( /* static */ absl::StatusOr ShapeInference::InferDotOpShape( const Shape& lhs, const Shape& rhs, const DotDimensionNumbers& dimension_numbers, - std::optional preferred_element_type) { + std::optional preferred_element_type, + const SparsityConfig& sparsity_config) { ABSL_RETURN_IF_ERROR(ExpectArray(lhs, "lhs of dot")); ABSL_RETURN_IF_ERROR(ExpectArray(rhs, "rhs of dot")); @@ -962,7 +992,8 @@ void GenerateDotResultDimensions( ABSL_RETURN_IF_ERROR(ValidateDotDimensionNumbers(lhs, rhs, dimension_numbers)); // Check the number and sizes of batch and contracting dimensions. - ABSL_RETURN_IF_ERROR(CheckDotDimensionConstraints(lhs, rhs, dimension_numbers)); + ABSL_RETURN_IF_ERROR(CheckDotDimensionConstraints(lhs, rhs, dimension_numbers, + sparsity_config)); std::vector dimensions; std::vector is_dynamic; @@ -2495,8 +2526,12 @@ absl::StatusOr InferWgradConvolveShape( for (int i = 0; i < num_spatial_dims; ++i) { input_spatial_dims[i] = lhs.dimensions(dnums.input_spatial_dimensions(i)); } - const int64_t input_features = - lhs.dimensions(dnums.input_feature_dimension()); + int64_t input_features = lhs.dimensions(dnums.input_feature_dimension()); + if (sparsity_config.has_lhs() && + sparsity_config.lhs().dimension() == dnums.input_feature_dimension()) { + ABSL_ASSIGN_OR_RETURN(input_features, + GetSparseDimSize(input_features, sparsity_config.lhs())); + } const int64_t input_batch = lhs.dimensions(dnums.input_batch_dimension()); std::vector kernel_spatial_dims(num_spatial_dims); @@ -2505,19 +2540,14 @@ absl::StatusOr InferWgradConvolveShape( } int64_t kernel_input_features = rhs.dimensions(dnums.kernel_input_feature_dimension()); - if (sparsity_config.has_rhs()) { + if (sparsity_config.has_rhs() && sparsity_config.rhs().dimension() == + dnums.kernel_input_feature_dimension()) { VLOG(8) << "Using sparse RHS for convolution. Got kernel_input_features: " << kernel_input_features << ", sparsity_config: " << SparsityConfigToString(sparsity_config); - int64_t num_non_zero = sparsity_config.rhs().num_non_zero(); - int64_t block_size = sparsity_config.rhs().block_size(); - if (num_non_zero != 1) { - return InvalidArgument("Only 1:N sparsity is currently supported."); - } - // Since the kernel is sparse, the effective number of input features is - // the number of non-zero elements times the block size. This currently - // assumes 1:N sparsity, where N is the block size. - kernel_input_features = kernel_input_features * block_size; + ABSL_ASSIGN_OR_RETURN( + kernel_input_features, + GetSparseDimSize(kernel_input_features, sparsity_config.rhs())); } else { VLOG(8) << "Not using sparse RHS for convolution."; } diff --git a/third_party/xla/xla/service/shape_inference.h b/third_party/xla/xla/service/shape_inference.h index c5566066296781..3fb1e9d0c657d3 100644 --- a/third_party/xla/xla/service/shape_inference.h +++ b/third_party/xla/xla/service/shape_inference.h @@ -374,7 +374,8 @@ class ShapeInference { static absl::StatusOr InferDotOpShape( const Shape& lhs, const Shape& rhs, const DotDimensionNumbers& dimension_numbers, - std::optional preferred_element_type); + std::optional preferred_element_type, + const SparsityConfig& sparsity_config = {}); // Helper that infers the shape produced by performing a ragged dot operation // with the given LHS and RHS shapes. An optional preferred_element_type can diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc index c8c75fa2673d26..4cbac0b5bbd7ce 100644 --- a/third_party/xla/xla/service/shape_inference_test.cc +++ b/third_party/xla/xla/service/shape_inference_test.cc @@ -3505,7 +3505,45 @@ TEST_F(ShapeInferenceTest, ConvWithSparsityFail) { /*preferred_element_type=*/std::nullopt); EXPECT_FALSE(status_or.ok()); EXPECT_THAT(status_or.status().message(), - HasSubstr("Only 1:N sparsity is currently supported.")); + HasSubstr("num_non_zero and block_size must be positive")); +} + +TEST_F(ShapeInferenceTest, ConvAndDotWithMtoNSparsity) { + const Shape lhs = ShapeUtil::MakeShape(F32, {256, 256}); + const Shape rhs = ShapeUtil::MakeShape(F32, {128, 256}); + SparsityConfig sp; + auto* rhs_sp = sp.mutable_rhs(); + rhs_sp->set_num_non_zero(2); + rhs_sp->set_block_size(4); + rhs_sp->set_dimension(0); + rhs_sp->set_stride(1); + + ConvolutionDimensionNumbers cdnums; + cdnums.set_input_batch_dimension(0); + cdnums.set_input_feature_dimension(1); + cdnums.set_kernel_input_feature_dimension(0); + cdnums.set_kernel_output_feature_dimension(1); + cdnums.set_output_batch_dimension(0); + cdnums.set_output_feature_dimension(1); + ASSERT_OK_AND_ASSIGN(Shape conv_shape, + ShapeInference::InferConvolveShape( + lhs, rhs, 1, 1, Window(), cdnums, sp, std::nullopt)); + EXPECT_TRUE(ShapeUtil::Equal(conv_shape, lhs)); + + DotDimensionNumbers ddnums; + ddnums.add_lhs_contracting_dimensions(1); + ddnums.add_rhs_contracting_dimensions(0); + ASSERT_OK_AND_ASSIGN( + Shape dot_shape, + ShapeInference::InferDotOpShape(lhs, rhs, ddnums, std::nullopt, sp)); + EXPECT_TRUE(ShapeUtil::Equal(dot_shape, lhs)); + + const Shape bad_rhs = ShapeUtil::MakeShape(F32, {127, 256}); + auto bad_status = + ShapeInference::InferDotOpShape(lhs, bad_rhs, ddnums, std::nullopt, sp); + EXPECT_FALSE(bad_status.ok()); + EXPECT_THAT(bad_status.status().message(), + HasSubstr("must be divisible by num_non_zero")); } TEST_F(ShapeInferenceTest, InferStochasticConvertShape) { diff --git a/third_party/xla/xla/service/spmd/dot_handler.cc b/third_party/xla/xla/service/spmd/dot_handler.cc index a91e86869f8eef..c1c5fabe94cef0 100644 --- a/third_party/xla/xla/service/spmd/dot_handler.cc +++ b/third_party/xla/xla/service/spmd/dot_handler.cc @@ -148,6 +148,10 @@ class CreateShardedDotFunctor final const Window&) const override { HloInstruction* l = ll.hlo(); HloInstruction* r = rr.hlo(); + CHECK(!dot_->sparsity_config().has_lhs() && + !dot_->sparsity_config().has_rhs()); + CHECK(!dot_->block_scaling_config().has_lhs() && + !dot_->block_scaling_config().has_rhs()); ABSL_ASSIGN_OR_RETURN( auto sharded_dot_shape, ShapeInference::InferDotOpShape( @@ -167,6 +171,15 @@ class CreateShardedDotFunctor final }; absl::Status SpmdPartitioningVisitor::HandleDot(HloInstruction* hlo) { + if (hlo->sharding().IsSingleDevice()) { + return DefaultAction(hlo); + } + // TODO(b/535773961): Support sharding for scaled / sparse dots. + if (hlo->block_scaling_config().has_lhs() || + hlo->block_scaling_config().has_rhs() || + hlo->sparsity_config().has_lhs() || hlo->sparsity_config().has_rhs()) { + return DefaultAction(hlo); + } if (!options_.need_resolve_conflicts && !options_.enable_windowed_einsum_for_all_gather && !options_.enable_windowed_einsum_for_reduce_scatter) { diff --git a/third_party/xla/xla/service/transpose_folding.cc b/third_party/xla/xla/service/transpose_folding.cc index 786b9418101916..9d63672018bf91 100644 --- a/third_party/xla/xla/service/transpose_folding.cc +++ b/third_party/xla/xla/service/transpose_folding.cc @@ -229,8 +229,9 @@ absl::StatusOr TransposeFolding::RunImpl( FunctionVisitor visit_fn([this, &foldable_dots, &foldable_convolutions]( HloInstruction* instruction) { if (instruction->opcode() == HloOpcode::kDot) { - // Don't fold dots with a 1D operand. - if ((instruction->operand(0)->shape().dimensions().size() < 2) || + // Don't fold dots with extra scaling/sparsity operands or a 1D operand. + if (instruction->operand_count() > 2 || + (instruction->operand(0)->shape().dimensions().size() < 2) || (instruction->operand(1)->shape().dimensions().size() < 2)) { return absl::OkStatus(); } diff --git a/third_party/xla/xla/service/transpose_folding_test.cc b/third_party/xla/xla/service/transpose_folding_test.cc index 79fc7854509830..a6cc72aeee6ef6 100644 --- a/third_party/xla/xla/service/transpose_folding_test.cc +++ b/third_party/xla/xla/service/transpose_folding_test.cc @@ -762,16 +762,18 @@ ENTRY entry_computation { op::Transpose(op::Parameter(2)))); } -TEST_F(TransposeFoldingTest, DoNotFoldConvWithNonSparsityAdditionalOperand) { +TEST_F(TransposeFoldingTest, DoNotFoldTransposeIntoBlockScalingConvOperand) { constexpr absl::string_view kHloString = R"( -HloModule DoNotFoldConvWithNonSparsityAdditionalOperand +HloModule DoNotFoldTransposeIntoBlockScalingConvOperand ENTRY entry_computation { input = f32[8,1] parameter(0) input_t = f32[1,8] transpose(input), dimensions={1,0} filter = f32[8,3] parameter(1) - extra = f32[] parameter(2) - ROOT conv = f32[1,3] convolution(input_t, filter, extra), dim_labels=bf_io->bf + extra = f32[1,1] parameter(2) + ROOT conv = f32[1,3] convolution(input_t, filter, extra), + dim_labels=bf_io->bf, + block_scaling_config={rhs={scale_idx=2 strides=1x4 steps=1x1}} } )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr module, From 7d5b9848da24acd3ded489467e471e383688d415 Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Sun, 27 Sep 2026 10:37:38 -0700 Subject: [PATCH 17/25] [XLA:CPU] Cap vector tile size and unroll large or batched vector.transpose ops PiperOrigin-RevId: 989230950 --- .../cpu/codegen/tiled/tiled_fusion_emitter.cc | 27 +++++++------------ .../transforms/tests/unroll_vectors.mlir | 22 +++++++++++++++ .../tiled/transforms/unroll_vectors_pass.cc | 25 ++++++++++++++--- 3 files changed, 54 insertions(+), 20 deletions(-) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tiled_fusion_emitter.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/tiled_fusion_emitter.cc index 004b6172639a31..1da7a36071bef9 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tiled_fusion_emitter.cc +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tiled_fusion_emitter.cc @@ -446,28 +446,21 @@ absl::StatusOr GetTiledHloComputation( // 2. Evaluate all candidates by substituting concrete tile sizes into the // symbolic tiles of roots and operands. - bool use_new_xtile_lowering = fusion.GetModule() - ->config() - .debug_options() - .xla_cpu_use_new_xtile_lowering(); - TargetMachineOptions target_machine_options( - fusion.GetModule()->config().debug_options()); + const DebugOptions& debug_options = + fusion.GetModule()->config().debug_options(); + bool use_new_xtile_lowering = debug_options.xla_cpu_use_new_xtile_lowering(); + TargetMachineOptions target_machine_options(debug_options); bool has_avx512 = absl::c_any_of( target_machine_options.enabled_features(), [](absl::string_view feature) { return absl::StrContains(feature, "avx512"); }); - int64_t max_bit_width = 8; - for (const HloInstruction* instr : - fusion.fused_instructions_computation()->instructions()) { - ShapeUtil::ForEachSubshape(instr->shape(), [&](const Shape& subshape, - const ShapeIndex& index) { - if (subshape.IsArray()) { - max_bit_width = std::max( - max_bit_width, primitive_util::BitWidth(subshape.element_type())); - } - }); + int64_t prefer_vector_width = debug_options.xla_cpu_prefer_vector_width(); + if (prefer_vector_width <= 0) { + prefer_vector_width = 256; } - int64_t max_vector_tile_size = (has_avx512 ? 512 : 256) / max_bit_width; + int64_t vector_width = + std::min(prefer_vector_width, has_avx512 ? 512 : 256); + int64_t max_vector_tile_size = vector_width / 32; struct Candidate { SmallVector padded_tile_sizes; diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/unroll_vectors.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/unroll_vectors.mlir index 11167b7253fcd7..33a4c420d179d1 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/unroll_vectors.mlir +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/unroll_vectors.mlir @@ -115,3 +115,25 @@ func.func @no_unroll_transpose(%arg0: vector<2x4xf32>) -> vector<4x2xf32> { // CHECK-LABEL: func.func @no_unroll_transpose( // CHECK: vector.transpose %{{.*}}, [1, 0] : vector<2x4xf32> to vector<4x2xf32> +// ----- + +func.func @unroll_large_transpose(%arg0: vector<32x16xf32>) + -> vector<16x32xf32> { + %0 = vector.transpose %arg0, [1, 0] + : vector<32x16xf32> to vector<16x32xf32> + return %0 : vector<16x32xf32> +} +// CHECK-LABEL: func.func @unroll_large_transpose( +// CHECK-COUNT-2: vector.transpose %{{.*}}, [1, 0] : vector<16x16xf32> to vector<16x16xf32> + +// ----- + +func.func @unroll_batch_transpose(%arg0: vector<2x4x8xf32>) + -> vector<2x8x4xf32> { + %0 = vector.transpose %arg0, [0, 2, 1] + : vector<2x4x8xf32> to vector<2x8x4xf32> + return %0 : vector<2x8x4xf32> +} +// CHECK-LABEL: func.func @unroll_batch_transpose( +// CHECK-COUNT-2: vector.transpose %{{.*}}, [0, 2, 1] : vector<1x4x8xf32> to vector<1x8x4xf32> + diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/unroll_vectors_pass.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/unroll_vectors_pass.cc index 35b390eb44e9e5..3d4bbe10a0bee7 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/unroll_vectors_pass.cc +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/unroll_vectors_pass.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -41,13 +42,23 @@ namespace xla::cpu { namespace { +llvm::SmallVector GetTransposeTargetShape( + mlir::vector::TransposeOp transpose_op, llvm::ArrayRef shape) { + llvm::ArrayRef perm = transpose_op.getPermutation(); + int64_t rank = shape.size(); + llvm::SmallVector target_shape(rank, 1); + for (int64_t i = 0; i < rank; ++i) { + if (i == rank - 1 || perm[i] == rank - 1) { + target_shape[i] = std::min(shape[i], 16); + } + } + return target_shape; +} + // Returns the target 1D unroll shape for a vector operation by keeping the // innermost dimension and reducing all leading dimensions to 1. static std::optional> GetNativeShape( mlir::Operation* op) { - if (mlir::isa(op)) { - return std::nullopt; - } auto unrollable = mlir::dyn_cast(op); if (!unrollable) { return std::nullopt; @@ -56,6 +67,14 @@ static std::optional> GetNativeShape( if (!shape || shape->size() <= 1) { return std::nullopt; } + if (auto transpose_op = mlir::dyn_cast(op)) { + llvm::SmallVector target_shape = + GetTransposeTargetShape(transpose_op, *shape); + if (target_shape == *shape) { + return std::nullopt; + } + return target_shape; + } if (llvm::all_of(llvm::ArrayRef(*shape).drop_back(), [](int64_t d) { return d == 1; })) { return std::nullopt; From 98b4ae0c0c57324ea3e956de96cd6415a249bbdd Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sun, 27 Sep 2026 12:46:50 -0700 Subject: [PATCH 18/25] Reverts f87634eaff1171e67c8d98cb04198cf168fc4c98 PiperOrigin-RevId: 989262062 --- .../xla/xla/hlo/utils/hlo_live_range.cc | 51 +---- .../xla/xla/hlo/utils/hlo_live_range_test.cc | 212 ------------------ 2 files changed, 2 insertions(+), 261 deletions(-) diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.cc b/third_party/xla/xla/hlo/utils/hlo_live_range.cc index b2dfdaca4e8661..f3027bb4cff609 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.cc @@ -368,29 +368,11 @@ void HloLiveRange::CalculateBufferStartEndMap() { instruction.IsRoot() ? computation_span_times_[computation].end : entry.second; - // If the instruction is in an asynchronous context, adjust its live range - // to cover the async window precisely: - // - end_time is extended to the async-done instruction so the buffer - // remains live for the full async duration. - // - start_time is tightened to the first-fully-bound instruction - // (async-start for standard chains, async-update for late-binding - // chains), because FlattenSchedule inlines the async computation - // immediately before that instruction and the inner buffer is not live - // before it fires. + // If the instruction is in an asynchronous context, extend the live range + // until the end of the async-done instruction. auto async_context_it = computations_in_async_context_.find(computation); if (async_context_it != computations_in_async_context_.end()) { const HloComputation* async_context = async_context_it->second; - // async_context can have multiple callers when the same async-wrapped - // computation is shared by several async-start/update instructions - // (e.g. one outside and one inside a while loop). FlattenSchedule only - // ever inlines the computation once, at whichever caller the schedule - // walk reaches first, so we take the minimum schedule time across all - // first-fully-bound callers below. This is deterministic (unlike - // picking an arbitrary caller from caller_instructions(), which is - // returned in no particular order) and never overshoots the schedule - // position the instructions were actually flattened at, which would - // otherwise push start_time past definition_end_time. - std::optional tightened_start_time; for (const HloInstruction* caller : async_context->caller_instructions()) { if (caller->IsAsynchronous()) { @@ -415,39 +397,10 @@ void HloLiveRange::CalculateBufferStartEndMap() { definition_end_time = std::max( definition_end_time, computation_span_times_[computation].end); } - // Track the earliest first-fully-bound caller's schedule time. - // FlattenSchedule inlines the async computation's instructions - // immediately before the first-fully-bound instruction (async-start - // for standard chains, async-update for late-binding chains) that - // actually triggers the flattening. The inner buffer is not live - // before that point. - absl::StatusOr is_first_fully_bound = - hlo_instruction_utils::async::IsFirstFullyBound(caller); - if (is_first_fully_bound.ok() && *is_first_fully_bound) { - auto first_bound_it = instruction_schedule_.find(caller); - if (first_bound_it != instruction_schedule_.end()) { - tightened_start_time = - tightened_start_time.has_value() - ? std::min(*tightened_start_time, first_bound_it->second) - : first_bound_it->second; - } - } } } - if (tightened_start_time.has_value()) { - // Cap by definition_end_time (computed above): for computations - // reached transitively through a kCall inside the async-wrapped - // computation, FlattenSchedule inlines their bodies even earlier - // than the async-wrapped computation's own instructions, so a - // caller's schedule time can land after such an instruction's own - // (already-correct) end. Capping avoids pushing start past end - // while still tightening whenever it's safe to do so. - start_time = std::min(*tightened_start_time, definition_end_time); - } VLOG(2) << "Setting the definition end time for op in async context: " << definition_end_time; - VLOG(2) << "Setting the definition start time for op in async context: " - << start_time; } for (const HloValue* value : diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc b/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc index c77a682e24a9d7..e3bb28c339bff7 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range_test.cc @@ -1202,217 +1202,5 @@ ENTRY entry { EXPECT_FALSE(root_buffers.empty()); EXPECT_TRUE(HloLiveRange::BufferLivesOut(*root_buffers[0], *aa, entry)); } - -TEST_F(HloLiveRangeTest, AsyncInnerBufferStartExtendedToOuterStart) { - const std::string hlo_string = R"hlo( - HloModule AsyncInnerBufferStart, is_scheduled=true - - %async_wrapped (p: f32[4]) -> f32[4] { - %p = f32[4] parameter(0) - ROOT %inner_op = f32[4] negate(%p) - } - - ENTRY %main (a: f32[4]) -> f32[4] { - %a = f32[4] parameter(0) - %outer_op = f32[4] negate(%a) - %async-start = ((f32[4]), f32[4], u32[]) async-start(%outer_op), - calls=%async_wrapped - %async-done = f32[4] async-done(%async-start) - ROOT %result = f32[4] add(%outer_op, %async-done) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(alias_analysis_, - HloAliasAnalysis::Run(module_.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(hlo_live_range_, - HloLiveRange::Run(module_->schedule(), *alias_analysis_, - module_->entry_computation())); - CheckSchedule(); - - HloComputation* async_wrapped = - module_->GetComputationWithName("async_wrapped"); - ASSERT_NE(async_wrapped, nullptr); - const HloInstruction* inner_op = - async_wrapped->GetInstructionWithName("inner_op"); - ASSERT_NE(inner_op, nullptr); - - const HloInstruction* async_done = - module_->entry_computation()->GetInstructionWithName("async-done"); - ASSERT_NE(async_done, nullptr); - const HloInstruction* result = - module_->entry_computation()->GetInstructionWithName("result"); - ASSERT_NE(result, nullptr); - - const HloInstruction* async_start = - module_->entry_computation()->GetInstructionWithName("async-start"); - ASSERT_NE(async_start, nullptr); - - auto inner_range = LiveRangeAt(inner_op); - auto async_start_time = - hlo_live_range_->instruction_schedule().at(async_start); - auto async_done_time = hlo_live_range_->instruction_schedule().at(async_done); - auto result_time = hlo_live_range_->instruction_schedule().at(result); - - EXPECT_EQ(inner_range.start, async_start_time); - EXPECT_GT(result_time, async_done_time); - EXPECT_EQ(inner_range.end, result_time); -} - -TEST_F(HloLiveRangeTest, AsyncInnerParameterStartExtendedToOuterStart) { - const std::string hlo_string = R"hlo( - HloModule AsyncInnerParamStart, is_scheduled=true - - %async_wrapped (p: f32[4]) -> f32[4] { - %p = f32[4] parameter(0) - ROOT %inner_op = f32[4] negate(%p) - } - - ENTRY %main (a: f32[4]) -> f32[4] { - %a = f32[4] parameter(0) - %async-start = ((f32[4]), f32[4], u32[]) async-start(%a), - calls=%async_wrapped - %async-done = f32[4] async-done(%async-start) - ROOT %result = f32[4] negate(%async-done) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(alias_analysis_, - HloAliasAnalysis::Run(module_.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(hlo_live_range_, - HloLiveRange::Run(module_->schedule(), *alias_analysis_, - module_->entry_computation())); - CheckSchedule(); - - HloComputation* async_wrapped = - module_->GetComputationWithName("async_wrapped"); - ASSERT_NE(async_wrapped, nullptr); - - const HloInstruction* async_start = - module_->entry_computation()->GetInstructionWithName("async-start"); - ASSERT_NE(async_start, nullptr); - const HloInstruction* inner_param = - async_wrapped->GetInstructionWithName("p"); - const HloInstruction* inner_op = - async_wrapped->GetInstructionWithName("inner_op"); - ASSERT_NE(inner_param, nullptr); - ASSERT_NE(inner_op, nullptr); - - auto async_start_time = - hlo_live_range_->instruction_schedule().at(async_start); - - EXPECT_EQ(LiveRangeAt(inner_param).start, 0); - EXPECT_EQ(LiveRangeAt(inner_op).start, async_start_time); -} - -TEST_F(HloLiveRangeTest, AsyncInnerBufferStartIsAsyncStartTime) { - const std::string hlo_string = R"hlo( - HloModule AsyncInnerOverlap, is_scheduled=true - - %async_wrapped (p: f32[4]) -> f32[4] { - %p = f32[4] parameter(0) - ROOT %inner_op = f32[4] negate(%p) - } - - ENTRY %main (a: f32[4]) -> f32[4] { - %a = f32[4] parameter(0) - %temp = f32[4] negate(%a) - %outer = f32[4] negate(%temp) - %async-start = ((f32[4]), f32[4], u32[]) async-start(%outer), - calls=%async_wrapped - %async-done = f32[4] async-done(%async-start) - ROOT %result = f32[4] negate(%async-done) - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(alias_analysis_, - HloAliasAnalysis::Run(module_.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(hlo_live_range_, - HloLiveRange::Run(module_->schedule(), *alias_analysis_, - module_->entry_computation())); - CheckSchedule(); - - const HloInstruction* async_start = - module_->entry_computation()->GetInstructionWithName("async-start"); - const HloInstruction* temp = - module_->entry_computation()->GetInstructionWithName("temp"); - HloComputation* async_wrapped = - module_->GetComputationWithName("async_wrapped"); - ASSERT_NE(async_start, nullptr); - ASSERT_NE(temp, nullptr); - ASSERT_NE(async_wrapped, nullptr); - - const HloInstruction* inner_op = - async_wrapped->GetInstructionWithName("inner_op"); - ASSERT_NE(inner_op, nullptr); - - auto async_start_time = - hlo_live_range_->instruction_schedule().at(async_start); - auto temp_range = LiveRangeAt(temp); - auto inner_range = LiveRangeAt(inner_op); - - EXPECT_EQ(inner_range.start, async_start_time); - EXPECT_GT(inner_range.start, temp_range.end); -} - -TEST_F(HloLiveRangeTest, AsyncComputationSharedByMultipleCallersUsesMinStart) { - const std::string hlo_string = R"hlo( - HloModule m, is_scheduled=true - - %nested (q: f32[4]) -> f32[4] { - %q = f32[4] parameter(0) - %t = f32[4] negate(%q) - ROOT %n = f32[4] negate(%t) - } - - %async_computation (p: f32[4]) -> f32[4] { - %p = f32[4] parameter(0) - ROOT %c = f32[4] call(%p), to_apply=%nested - } - - %while_cond (wp: ((f32[4]), f32[4], u32[])) -> pred[] { - %wp = ((f32[4]), f32[4], u32[]) parameter(0) - ROOT %cond = pred[] constant(true) - } - - %while_body (bp: ((f32[4]), f32[4], u32[])) -> ((f32[4]), f32[4], u32[]) { - %bp = ((f32[4]), f32[4], u32[]) parameter(0) - %g = f32[4] get-tuple-element(%bp), index=1 - ROOT %s = ((f32[4]), f32[4], u32[]) async-start(%g), calls=%async_computation - } - - ENTRY %main (a: f32[4]) -> ((f32[4]), f32[4], u32[]) { - %a = f32[4] parameter(0) - %init = ((f32[4]), f32[4], u32[]) async-start(%a), calls=%async_computation - ROOT %loop = ((f32[4]), f32[4], u32[]) while(%init), condition=%while_cond, body=%while_body - } - )hlo"; - - ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); - ASSERT_OK_AND_ASSIGN(alias_analysis_, - HloAliasAnalysis::Run(module_.get(), &alias_info_)); - ASSERT_OK_AND_ASSIGN(hlo_live_range_, - HloLiveRange::Run(module_->schedule(), *alias_analysis_, - module_->entry_computation())); - CheckSchedule(); - - HloComputation* nested = module_->GetComputationWithName("nested"); - ASSERT_NE(nested, nullptr); - const HloInstruction* t = nested->GetInstructionWithName("t"); - ASSERT_NE(t, nullptr); - - const HloInstruction* init = - module_->entry_computation()->GetInstructionWithName("init"); - ASSERT_NE(init, nullptr); - - auto init_time = hlo_live_range_->instruction_schedule().at(init); - auto t_range = LiveRangeAt(t); - - EXPECT_LE(t_range.start, t_range.end); - EXPECT_LT(t_range.start, init_time); -} - } // namespace } // namespace xla From 66585ca45169805cf38dbf6f585c23b80510bc3c Mon Sep 17 00:00:00 2001 From: Alexander Belyaev Date: Sun, 27 Sep 2026 12:48:03 -0700 Subject: [PATCH 19/25] [XLA] Cache xla.range for apply_indexing op. Currently, we remove the constraints from the indexing maps before constructing ApplyIndexingOps. Because of that we miss some of the arithmetic simplifications, because with the contraints the result range might be smaller and that would in turn enable some simplification/folding. I think a proper fix would be to allow constraints in apply_indexing. However, that could affect the compile time severely. I will experiment with that in a follow up. Even if we allow constraints, it makes sense to cache the ranges to save on compute a bit. PiperOrigin-RevId: 989262341 --- .../tests/elementwise/minimum_2d_tile.hlo | 15 +++--- .../minimum_three_tiles_per_block.hlo | 2 +- third_party/xla/xla/codegen/emitters/BUILD | 1 + .../codegen/emitters/elemental_hlo_to_mlir.cc | 29 ++++++++++- .../xla/xla/codegen/emitters/ir/xla_ops.cc | 51 ++++++++++++------- .../xla/xla/codegen/emitters/ir/xla_ops.h | 9 ++++ .../tests/loop/pad_slice_bounds_check.hlo | 19 +++++++ .../emitters/transforms/simplify_arith.cc | 4 +- .../transforms/tests/simplify_arith.mlir | 15 +++++- .../xla/xla/hlo/analysis/indexing_map.cc | 10 ++++ .../xla/xla/hlo/analysis/indexing_map.h | 3 ++ .../xla/xla/hlo/analysis/indexing_map_test.cc | 16 ++++++ third_party/xla/xla/hlo/analysis/interval.cc | 29 +++++------ third_party/xla/xla/hlo/analysis/interval.h | 1 + 14 files changed, 156 insertions(+), 48 deletions(-) create mode 100644 third_party/xla/xla/codegen/emitters/tests/loop/pad_slice_bounds_check.hlo diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo index 44bb8b50bb6a91..db738571de07f8 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo @@ -19,22 +19,23 @@ fusion { p1 = f32[128,256] parameter(1) ROOT minimum = f32[128,256] minimum(p0, p1) } -// CHECK: #indexing_map = #xla.indexing_map<"(pid, tid) -> ((tid / 64) * 8 + pid * 64), domain: pid in [0, 1], tid in [0, 511]"> -// CHECK: #indexing_map1 = #xla.indexing_map<"(tid) -> ((tid mod 64) * 4), domain: tid in [0, 511]"> +// CHECK: #indexing_map = #xla.indexing_map<"(pid) -> (min(pid * -86 + 1024, 86)), domain: pid in [0, 11]"> +// CHECK: #indexing_map1 = #xla.indexing_map<"(pid, tid) -> (((pid * 86 + tid) / 64) * 8), domain: pid in [0, 11], tid in [0, 85]"> +// CHECK: #indexing_map2 = #xla.indexing_map<"(pid, tid) -> (((pid * 86 + tid) mod 64) * 4), domain: pid in [0, 11], tid in [0, 85]"> // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 1 : index]}) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 11 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index -// CHECK: %[[LIMIT:.*]] = arith.constant 512 : index +// CHECK: %[[LIMIT:.*]] = xla.apply_indexing #indexing_map(%[[PID]]) // CHECK: scf.for %[[TID:.*]] = %[[C0]] to %[[LIMIT]] step %[[C1]] { -// CHECK: %[[IDX_0:.*]] = xla.apply_indexing #indexing_map(%[[PID]], %[[TID]]) -// CHECK: %[[IDX_1:.*]] = xla.apply_indexing #indexing_map1(%[[TID]]) +// CHECK: %[[IDX_0:.*]] = xla.apply_indexing #indexing_map1(%[[PID]], %[[TID]]) {xla.range = [0 : index, 120 : index]} +// CHECK: %[[IDX_1:.*]] = xla.apply_indexing #indexing_map2(%[[PID]], %[[TID]]) // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]][%[[IDX_0]], %[[IDX_1]]] [8, 4] [1, 1] // CHECK: %[[T_1:.*]] = xtile.extract %[[RHS]][%[[IDX_0]], %[[IDX_1]]] [8, 4] [1, 1] // CHECK: %[[MIN:.*]] = stablehlo.minimum %[[T_0]], %[[T_1]] : tensor<8x4xf32> @@ -73,7 +74,7 @@ ENTRY main { "fusion_backend_config":{ "block_level_fusion_config":{ "output_tiles":[{"sizes":["8", "4"]}], - "num_ctas":"2" + "num_ctas":"12" } } } diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo index 6fbe52e11dfa41..513ac7a83b328d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo @@ -37,7 +37,7 @@ ENTRY main { // CHECK: %[[UB:.*]] = xla.apply_indexing #[[$UB_MAP]](%[[PID]]) // CHECK: scf.for %[[TID:.*]] = %[[C0]] to %[[UB]] step %[[C1]] { -// CHECK: %[[GLOB_TID:.*]] = xla.apply_indexing #[[$TILE_ID_MAP]](%[[PID]], %[[TID]]) +// CHECK: %[[GLOB_TID:.*]] = xla.apply_indexing #[[$TILE_ID_MAP]](%[[PID]], %[[TID]]) {xla.range = [0 : index, 7 : index]} // CHECK: %[[LHS_TILE:.*]] = xtile.extract %[[LHS]][%[[GLOB_TID]], %[[C0]]] [1, 4] [1, 1] : memref<8x4xf32> -> tensor<1x4xf32> // CHECK: %[[RHS_TILE:.*]] = xtile.extract %[[RHS]][%[[GLOB_TID]], %[[C0]]] [1, 4] [1, 1] : memref<8x4xf32> -> tensor<1x4xf32> // CHECK: %[[MIN:.*]] = arith.minimumf %[[LHS_TILE]], %[[RHS_TILE]] : tensor<1x4xf32> diff --git a/third_party/xla/xla/codegen/emitters/BUILD b/third_party/xla/xla/codegen/emitters/BUILD index 0e987b5d9a85fc..6bee1b46cd61aa 100644 --- a/third_party/xla/xla/codegen/emitters/BUILD +++ b/third_party/xla/xla/codegen/emitters/BUILD @@ -108,6 +108,7 @@ cc_library( "//xla/codegen/emitters/ir:xla", "//xla/codegen/emitters/transforms:lowering_utils", "//xla/hlo/analysis:indexing_analysis", + "//xla/hlo/analysis:interval", "//xla/hlo/analysis:symbolic_map", "//xla/hlo/ir:hlo", "//xla/hlo/transforms/simplifiers:gather_simplifier", diff --git a/third_party/xla/xla/codegen/emitters/elemental_hlo_to_mlir.cc b/third_party/xla/xla/codegen/emitters/elemental_hlo_to_mlir.cc index 28ce0d0d80e566..40f9753de3bc02 100644 --- a/third_party/xla/xla/codegen/emitters/elemental_hlo_to_mlir.cc +++ b/third_party/xla/xla/codegen/emitters/elemental_hlo_to_mlir.cc @@ -72,6 +72,7 @@ limitations under the License. #include "xla/comparison_util.h" #include "xla/hlo/analysis/indexing_analysis.h" #include "xla/hlo/analysis/indexing_map.h" +#include "xla/hlo/analysis/interval.h" #include "xla/hlo/analysis/symbolic_expr.h" #include "xla/hlo/analysis/symbolic_map.h" #include "xla/hlo/ir/hlo_casting_utils.h" @@ -660,13 +661,37 @@ SmallVector MapElementwiseOp( SmallVector ApplyIndexing(IndexingMap map, ValueRange dims, ValueRange symbols, ImplicitLocOpBuilder& b) { - map.ClearConstraints(); SmallVector results; for (unsigned int i = 0; i < map.GetNumResults(); ++i) { SmallVector result; - b.createOrFold(result, dims, symbols, map.GetSubMap(i)); + IndexingMap sub_map = map.GetSubMap(i); + sub_map.ClearConstraints(); + b.createOrFold(result, dims, symbols, std::move(sub_map)); results.append(result); } + if (map.GetConstraintsCount() == 0) { + return results; + } + // Add constraints to the apply indexing ops. + // TODO(b/542571968): A more principled, but potentially a much more expensive + // way to fix this is to allow the constraints in the indexing maps for + // apply_indexing ops. That will require to update a lot of tests, but it is + // worth trying. + SmallVector result_ranges = map.ComputeResultRanges(); + for (const auto& [result, range] : llvm::zip(results, result_ranges)) { + // Bare dim/symbol results fold to a pre-existing operand. Its defining op + // may not be guarded by `map`'s constraints, so never annotate it. + if (llvm::is_contained(dims, result) || + llvm::is_contained(symbols, result)) { + continue; + } + auto apply_op = result.getDefiningOp(); + if (!apply_op || range.IsUnconstrained()) { + continue; + } + apply_op->setAttr("xla.range", + b.getIndexArrayAttr({range.lower, range.upper})); + } return results; } diff --git a/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc b/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc index 2def844cfcf441..d11b937d781f28 100644 --- a/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc +++ b/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc @@ -105,9 +105,14 @@ std::optional GetRange(mlir::Value value) { }; if (auto apply = value.getDefiningOp()) { - return apply.getIndexingMap().GetRangeEvaluator().ComputeExpressionRange( - apply.getIndexingMap().GetSymbolicMap().GetResult( - mlir::cast(value).getResultNumber())); + Interval range = + apply.getIndexingMap().GetRangeEvaluator().ComputeExpressionRange( + apply.getIndexingMap().GetSymbolicMap().GetResult( + mlir::cast(value).getResultNumber())); + if (auto attr_range = attr_to_range(apply->getAttr("xla.range"))) { + range = range.Intersect(*attr_range); + } + return range; } else if (auto cst = value.getDefiningOp()) { return {{cst.value(), cst.value()}}; } else if (value.getDefiningOp()) { @@ -352,6 +357,17 @@ absl::StatusOr GetNewIndexingMapAfterFoldingSequence( } } // namespace +void ReplaceApplyIndexingOp(PatternRewriter& rewriter, + ApplyIndexingOp indexing_op, ValueRange operands, + const IndexingMap& indexing_map) { + auto new_op = ApplyIndexingOp::create(rewriter, indexing_op.getLoc(), + operands, indexing_map); + if (auto range_attr = indexing_op->getAttr("xla.range")) { + new_op->setAttr("xla.range", range_attr); + } + rewriter.replaceOp(indexing_op, new_op.getResults()); +} + namespace { // Simplifies the indexing map. @@ -365,8 +381,8 @@ struct SimplifyIndexingMap : public mlir::OpRewritePattern { return rewriter.notifyMatchFailure(indexing_op, "IndexingMap is already simplified"); } - rewriter.replaceOpWithNewOp( - indexing_op, indexing_op.getOperands(), indexing_map); + ReplaceApplyIndexingOp(rewriter, indexing_op, indexing_op.getOperands(), + indexing_map); return success(); } }; @@ -390,8 +406,7 @@ struct RemoveUnusedVariables : public mlir::OpRewritePattern { operands.push_back(indexing_op.getOperand(i)); } } - rewriter.replaceOpWithNewOp(indexing_op, operands, - indexing_map); + ReplaceApplyIndexingOp(rewriter, indexing_op, operands, indexing_map); return success(); } }; @@ -405,9 +420,8 @@ struct MoveSymbolsToDims : public mlir::OpRewritePattern { if (indexing_map.GetSymbolCount() == 0) { return rewriter.notifyMatchFailure(indexing_op, "No symbols found"); } - rewriter.replaceOpWithNewOp( - indexing_op, indexing_op->getOperands(), - indexing_map.ConvertSymbolsToDimensions()); + ReplaceApplyIndexingOp(rewriter, indexing_op, indexing_op->getOperands(), + indexing_map.ConvertSymbolsToDimensions()); return success(); } }; @@ -475,8 +489,8 @@ struct FoldApplyIndexingSequence new_operands.append(replacement->added_dim_args); new_operands.append(begin + num_dims, begin + num_dims + num_syms); - rewriter.replaceOpWithNewOp(indexing_op, new_operands, - replacement->indexing_map); + ReplaceApplyIndexingOp(rewriter, indexing_op, new_operands, + replacement->indexing_map); return success(); } @@ -518,7 +532,7 @@ struct FoldApplyIndexingOperands unsigned new_num_operands = indexing_op->getNumOperands() - num_constants; SmallVector new_operands; new_operands.reserve(new_num_operands); - SmallVector new_dim_vars; + std::vector new_dim_vars; new_dim_vars.reserve(num_dims); unsigned new_num_dims = 0; @@ -540,11 +554,12 @@ struct FoldApplyIndexingOperands new_dim_vars.push_back(indexing_map.GetDimVar(operand_id)); } } - rewriter.replaceOpWithNewOp( - indexing_op, new_operands, - symbolic_map.ReplaceDimsAndSymbols(dim_replacements, {}, new_num_dims, - 0), - new_dim_vars, SmallVector()); + IndexingMap new_indexing_map(symbolic_map.ReplaceDimsAndSymbols( + dim_replacements, {}, new_num_dims, 0), + std::move(new_dim_vars), /*range_vars=*/{}, + /*rt_vars=*/{}); + ReplaceApplyIndexingOp(rewriter, indexing_op, new_operands, + new_indexing_map); return success(); } }; diff --git a/third_party/xla/xla/codegen/emitters/ir/xla_ops.h b/third_party/xla/xla/codegen/emitters/ir/xla_ops.h index a94c09cc5bc98f..e80a75d180e346 100644 --- a/third_party/xla/xla/codegen/emitters/ir/xla_ops.h +++ b/third_party/xla/xla/codegen/emitters/ir/xla_ops.h @@ -28,6 +28,8 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" // IWYU pragma: keep #include "mlir/IR/OpDefinition.h" // IWYU pragma: keep #include "mlir/IR/OpImplementation.h" // IWYU pragma: keep +#include "mlir/IR/PatternMatch.h" // IWYU pragma: keep +#include "mlir/IR/ValueRange.h" // IWYU pragma: keep #include "mlir/Interfaces/CallInterfaces.h" // IWYU pragma: keep #include "mlir/Interfaces/InferTypeOpInterface.h" // IWYU pragma: keep #include "mlir/Interfaces/SideEffectInterfaces.h" // IWYU pragma: keep @@ -78,6 +80,13 @@ std::optional GetBackendKind(mlir::func::FuncOp fn); void SetBackendKind(mlir::MLIRContext* context, mlir::func::FuncOp fn, xla::BackendKind backend_kind); +// Replaces `indexing_op` with a new `ApplyIndexingOp` with the given operands +// and indexing map, preserving the `xla.range` attribute if present. +void ReplaceApplyIndexingOp(mlir::PatternRewriter& rewriter, + ApplyIndexingOp indexing_op, + mlir::ValueRange operands, + const IndexingMap& indexing_map); + } // namespace xla #endif // XLA_CODEGEN_EMITTERS_IR_XLA_OPS_H_ diff --git a/third_party/xla/xla/codegen/emitters/tests/loop/pad_slice_bounds_check.hlo b/third_party/xla/xla/codegen/emitters/tests/loop/pad_slice_bounds_check.hlo new file mode 100644 index 00000000000000..d5a48c7773eae7 --- /dev/null +++ b/third_party/xla/xla/codegen/emitters/tests/loop/pad_slice_bounds_check.hlo @@ -0,0 +1,19 @@ +// RUN: gpu_fusion_to_mlir %s | emitters_opt --xla-test-optimize | FileCheck %s +// RUN: cpu_fusion_to_mlir %s | emitters_opt --xla-test-optimize | FileCheck %s +// RUN: gpu_test_correctness %s +// RUN: cpu_test_correctness %s + +fusion { + p0 = f32[4,4] parameter(0) + c0 = f32[] constant(0) + pad = f32[7,6] pad(p0, c0), padding=0_0_1x0_2 + ROOT slice = f32[7,5] slice(pad), slice={[0:7], [1:6]} +} + +// CHECK-LABEL: func.func private @fusion_slice( +// CHECK-SAME: %[[IN:[a-zA-Z0-9]*]]: tensor<4x4xf32> +// CHECK-SAME: %[[J:[a-zA-Z0-9]*]]: index {xla.range = [0 : index, 4 : index]} +// CHECK: %[[COL:.*]] = xla.apply_indexing #{{.*}}(%[[J]]) +// CHECK: arith.cmpi sle, %[[COL]], %{{.*}} : index +// CHECK: scf.if +// CHECK: tensor.extract %[[IN]][%{{.*}}, %[[COL]]] diff --git a/third_party/xla/xla/codegen/emitters/transforms/simplify_arith.cc b/third_party/xla/xla/codegen/emitters/transforms/simplify_arith.cc index 719bbc5ba68ac6..d59ba6cc590b38 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/simplify_arith.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/simplify_arith.cc @@ -429,8 +429,8 @@ struct RefineConstraints : public OpRewritePattern { return rewriter.notifyMatchFailure(indexing_op, "No bounds to refine"); } indexing_map.Simplify(); - rewriter.replaceOpWithNewOp( - indexing_op, indexing_op.getOperands(), indexing_map); + ReplaceApplyIndexingOp(rewriter, indexing_op, indexing_op.getOperands(), + indexing_map); return mlir::success(); } }; diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_arith.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_arith.mlir index 340005d916c7b1..39ba54845b1de6 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_arith.mlir +++ b/third_party/xla/xla/codegen/emitters/transforms/tests/simplify_arith.mlir @@ -479,7 +479,18 @@ module { // CHECK-NEXT: %[[CST:.*]] = arith.constant dense : tensor<4xi1> // CHECK-NEXT: return %[[CST]] +// ----- +#constrained_map = #xla.indexing_map<"(d0, d1) -> ((d0 * 21846 + d1) floordiv 256), domain: d0 in [0, 11], d1 in [0, 21845]"> +module { + func.func @apply_indexing_range_attr(%pid: index, %tid: index) -> i1 { + %row = xla.apply_indexing #constrained_map(%pid, %tid) {xla.range = [0 : index, 1023 : index]} + %c1024 = arith.constant 1024 : index + %in_bounds = arith.cmpi slt, %row, %c1024 : index + return %in_bounds : i1 + } +} - - +// CHECK-LABEL: @apply_indexing_range_attr +// CHECK-NEXT: %[[TRUE:.*]] = arith.constant true +// CHECK-NEXT: return %[[TRUE]] diff --git a/third_party/xla/xla/hlo/analysis/indexing_map.cc b/third_party/xla/xla/hlo/analysis/indexing_map.cc index 0823f425b59eb3..856499b94cf052 100644 --- a/third_party/xla/xla/hlo/analysis/indexing_map.cc +++ b/third_party/xla/xla/hlo/analysis/indexing_map.cc @@ -922,6 +922,16 @@ RangeEvaluator IndexingMap::GetRangeEvaluator() const { return RangeEvaluator(*this, GetMLIRContext()); } +llvm::SmallVector IndexingMap::ComputeResultRanges() const { + RangeEvaluator range_evaluator = GetRangeEvaluator(); + llvm::SmallVector ranges; + ranges.reserve(GetNumResults()); + for (SymbolicExpr expr : symbolic_map_.GetResults()) { + ranges.push_back(range_evaluator.ComputeExpressionRange(expr)); + } + return ranges; +} + const Interval& IndexingMap::GetDimensionBound(int64_t dim_id) const { return dim_vars_[dim_id].bounds; } diff --git a/third_party/xla/xla/hlo/analysis/indexing_map.h b/third_party/xla/xla/hlo/analysis/indexing_map.h index f4a0d85dee3276..3682f4061e17d2 100644 --- a/third_party/xla/xla/hlo/analysis/indexing_map.h +++ b/third_party/xla/xla/hlo/analysis/indexing_map.h @@ -173,6 +173,9 @@ class IndexingMap { // Returns the range evaluator for the indexing map's domain. RangeEvaluator GetRangeEvaluator() const; + // Computes the range for each result of the indexing map. + llvm::SmallVector ComputeResultRanges() const; + // Getters for dimension vars. const Variable& GetDimVar(int64_t id) const { return dim_vars_[id]; } const std::vector& GetDimVars() const { return dim_vars_; } diff --git a/third_party/xla/xla/hlo/analysis/indexing_map_test.cc b/third_party/xla/xla/hlo/analysis/indexing_map_test.cc index 386aabd54ff197..f9c917215e4bdd 100644 --- a/third_party/xla/xla/hlo/analysis/indexing_map_test.cc +++ b/third_party/xla/xla/hlo/analysis/indexing_map_test.cc @@ -1875,5 +1875,21 @@ TEST_F(IndexingMapTest, IsUndefined) { EXPECT_FALSE(defined_map.IsUndefined()); } +TEST_F(IndexingMapTest, ComputeResultRanges) { + IndexingMap indexing_map = Parse(R"( + (d0, d1) -> ((d0 * 86 + d1) floordiv 64, (d0 * 86 + d1) mod 64), + domain: + d0 in [0, 11], + d1 in [0, 85], + d0 * 86 + d1 in [0, 1023] + )"); + EXPECT_THAT(indexing_map.ComputeResultRanges(), + ElementsAre(Interval{0, 15}, Interval{0, 63})); + + indexing_map.ClearConstraints(); + EXPECT_THAT(indexing_map.ComputeResultRanges(), + ElementsAre(Interval{0, 16}, Interval{0, 63})); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/analysis/interval.cc b/third_party/xla/xla/hlo/analysis/interval.cc index 42a708ee5e6cd2..77c4e691051fca 100644 --- a/third_party/xla/xla/hlo/analysis/interval.cc +++ b/third_party/xla/xla/hlo/analysis/interval.cc @@ -29,6 +29,12 @@ limitations under the License. #include "llvm/Support/raw_ostream.h" namespace xla { +namespace { + +constexpr int64_t kMin = std::numeric_limits::min(); +constexpr int64_t kMax = std::numeric_limits::max(); + +} // namespace std::ostream& operator<<(std::ostream& out, const Interval& interval) { out << absl::StrFormat("[%d, %d]", interval.lower, interval.upper); @@ -39,6 +45,10 @@ std::string Interval::ToString() const { return absl::StrFormat("[%d, %d]", lower, upper); } +bool Interval::IsUnconstrained() const { + return lower == kMin && upper == kMax; +} + inline llvm::raw_ostream& operator<<(llvm::raw_ostream& os, const Interval& interval) { os << absl::StrFormat("[%d, %d]", interval.lower, interval.upper); @@ -49,8 +59,7 @@ int64_t Interval::GetLoopTripCount() const { if (!IsFeasible()) { return 0; } - DCHECK((static_cast(upper) - lower + 1) <= - std::numeric_limits::max()); + DCHECK((static_cast(upper) - lower + 1) <= kMax); return upper - lower + 1; } @@ -82,9 +91,6 @@ Interval Interval::operator+(const Interval& rhs) const { int64_t out_lower; int64_t out_upper; - constexpr int64_t kMin = std::numeric_limits::min(); - constexpr int64_t kMax = std::numeric_limits::max(); - bool lower_overflow = llvm::AddOverflow(lower, rhs.lower, out_lower); bool upper_overflow = llvm::AddOverflow(upper, rhs.upper, out_upper); @@ -110,9 +116,6 @@ Interval Interval::operator+(const Interval& rhs) const { } Interval Interval::operator*(const Interval& rhs) const { - constexpr int64_t kMin = std::numeric_limits::min(); - constexpr int64_t kMax = std::numeric_limits::max(); - auto mul = [&](int64_t p) { int64_t l = lower; int64_t u = upper; @@ -137,19 +140,13 @@ Interval Interval::operator*(const Interval& rhs) const { } Interval Interval::operator-() const { - int64_t ub = lower == std::numeric_limits::min() - ? std::numeric_limits::max() - : -lower; - int64_t lb = upper == std::numeric_limits::max() - ? std::numeric_limits::min() - : -upper; + int64_t ub = lower == kMin ? kMax : -lower; + int64_t lb = upper == kMax ? kMin : -upper; return Interval{lb, ub}; } Interval Interval::FloorDiv(int64_t rhs) const { auto saturate_div = [](int64_t lhs, int64_t rhs) { - constexpr int64_t kMin = std::numeric_limits::min(); - constexpr int64_t kMax = std::numeric_limits::max(); if (lhs == kMin) { return rhs > 0 ? kMin : kMax; } diff --git a/third_party/xla/xla/hlo/analysis/interval.h b/third_party/xla/xla/hlo/analysis/interval.h index e095460de0862b..a76a0e75b324c6 100644 --- a/third_party/xla/xla/hlo/analysis/interval.h +++ b/third_party/xla/xla/hlo/analysis/interval.h @@ -33,6 +33,7 @@ struct Interval { std::string ToString() const; bool IsPoint() const { return lower == upper; } bool IsFeasible() const { return lower <= upper; } + bool IsUnconstrained() const; // Returns the number of elements in the interval. Asserts that the number of // elements fits in an int64_t. For this reason, this should only be used for From 3df6868daf59b7805818f0a21dd94fc54540f1e9 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sun, 27 Sep 2026 20:29:21 -0700 Subject: [PATCH 20/25] Optimize tflite::reference_ops::StridedSlice via contiguous dimension coalescing and WriteN. PiperOrigin-RevId: 989374464 --- tensorflow/lite/kernels/BUILD | 7 + .../internal/reference/strided_slice.h | 68 ++- tensorflow/lite/kernels/strided_slice_test.cc | 566 ++++++++++++++++++ 3 files changed, 632 insertions(+), 9 deletions(-) diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD index d43aaa4fcbbdd6..855fdae1fd189b 100644 --- a/tensorflow/lite/kernels/BUILD +++ b/tensorflow/lite/kernels/BUILD @@ -2527,10 +2527,17 @@ cc_test( deps = [ ":test_main", ":test_util", + "//tensorflow/lite:string_util", + "//tensorflow/lite/core/c:common", + "//tensorflow/lite/kernels/internal:reference_base", + "//tensorflow/lite/kernels/internal:tensor", + "//tensorflow/lite/kernels/internal:types", "//tensorflow/lite/schema:schema_fbs", "//tensorflow/lite/types:half", + "@com_google_benchmark//:benchmark", "@com_google_googletest//:gtest", "@eigen_archive//:eigen3", + "@ruy//ruy/profiler:instrumentation", ], ) diff --git a/tensorflow/lite/kernels/internal/reference/strided_slice.h b/tensorflow/lite/kernels/internal/reference/strided_slice.h index 7307c47e00ab0b..d0f628b0f46a90 100644 --- a/tensorflow/lite/kernels/internal/reference/strided_slice.h +++ b/tensorflow/lite/kernels/internal/reference/strided_slice.h @@ -15,7 +15,10 @@ limitations under the License. #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_STRIDED_SLICE_H_ #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_STRIDED_SLICE_H_ -#include +#include +#include +#include +#include #include #include "ruy/profiler/instrumentation.h" // from @ruy @@ -108,27 +111,74 @@ inline void StridedSlice(const DynamicStridedSliceParams& op_params, for (int i = dims - 2; i >= 0; --i) { input_strides[i] = input_strides[i + 1] * input_shape.Dims(i + 1); } + auto loop_condition = [](int64_t index, int64_t stop, int stride) { + return stride > 0 ? index < stop : index > stop; + }; for (int axis = 0; axis < dims; ++axis) { starts[axis] = StartForAxis(op_params, input_shape, axis); stops[axis] = EndForAxis(op_params, input_shape, axis, starts[axis]); + if (!loop_condition(starts[axis], stops[axis], op_params.strides[axis])) { + return; + } } - auto loop_condition = [](int64_t index, int64_t stop, int stride) { - return stride > 0 ? index < stop : index > stop; + int inner_contig_axis = dims; + while ( + inner_contig_axis > 1 && op_params.strides[inner_contig_axis - 1] == 1 && + starts[inner_contig_axis - 1] == 0 && + stops[inner_contig_axis - 1] == input_shape.Dims(inner_contig_axis - 1)) { + --inner_contig_axis; + } + const int last_loop_axis = inner_contig_axis - 1; + const bool last_loop_stride_is_1 = op_params.strides[last_loop_axis] == 1; + const int64_t last_loop_contig_len = + last_loop_stride_is_1 ? (static_cast(stops[last_loop_axis]) - + static_cast(starts[last_loop_axis])) * + input_strides[last_loop_axis] + : 0; + + auto write_contiguous = [&](int64_t pos, int64_t count) { + if constexpr (std::is_trivially_copyable_v) { + if (count > 1) { + constexpr int64_t kMaxChunk = + std::numeric_limits::max() / static_cast(sizeof(T)); + while (count > 0) { + const int chunk = + static_cast(std::min(count, kMaxChunk)); + writer->WriteN(pos, chunk); + pos += chunk; + count -= chunk; + } + return; + } + } + for (int64_t i = 0; i < count; ++i) { + writer->Write(pos + i); + } }; - std::function write_slice = [&](int axis, - int64_t input_index) { - if (axis == dims) { - writer->Write(input_index); + + auto write_slice = [&](auto& self, int axis, int64_t input_index) -> void { + if (axis == last_loop_axis) { + if (last_loop_stride_is_1) { + write_contiguous(input_index + starts[axis] * input_strides[axis], + last_loop_contig_len); + } else { + for (int64_t offset = starts[axis]; + loop_condition(offset, stops[axis], op_params.strides[axis]); + offset += op_params.strides[axis]) { + write_contiguous(input_index + offset * input_strides[axis], + input_strides[axis]); + } + } return; } for (int64_t offset = starts[axis]; loop_condition(offset, stops[axis], op_params.strides[axis]); offset += op_params.strides[axis]) { - write_slice(axis + 1, input_index + offset * input_strides[axis]); + self(self, axis + 1, input_index + offset * input_strides[axis]); } }; - write_slice(/*axis=*/0, /*input_index=*/0); + write_slice(write_slice, /*axis=*/0, /*input_index=*/0); } template diff --git a/tensorflow/lite/kernels/strided_slice_test.cc b/tensorflow/lite/kernels/strided_slice_test.cc index b6dd4cd4958054..e2ec4d178d8738 100644 --- a/tensorflow/lite/kernels/strided_slice_test.cc +++ b/tensorflow/lite/kernels/strided_slice_test.cc @@ -12,9 +12,16 @@ 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/kernels/internal/reference/strided_slice.h" + #include +#include +#include +#include +#include #include +#include #include #include #include @@ -23,9 +30,16 @@ limitations under the License. #include #include +#include "benchmark/benchmark.h" // from @com_google_benchmark #include "Eigen/Core" // from @eigen_archive // IWYU pragma: keep +#include "ruy/profiler/instrumentation.h" // from @ruy +#include "tensorflow/lite/core/c/common.h" +#include "tensorflow/lite/kernels/internal/portable_tensor.h" +#include "tensorflow/lite/kernels/internal/runtime_shape.h" +#include "tensorflow/lite/kernels/internal/tensor.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" +#include "tensorflow/lite/string_util.h" #include "tensorflow/lite/types/half.h" namespace tflite { @@ -1631,5 +1645,557 @@ TYPED_TEST(StridedSliceOpTest, NoopOffset) { EXPECT_THAT(m.GetOutput(), ElementsAreTypedArray( CastVector({3, 2, 1, 6, 5, 4}))); } + +// Verbatim copy of the legacy element-by-element std::function recursive +// StridedSlice implementation used as a reference oracle to verify 100% +// bit-for-bit identical output across all tensor shapes, strides, and masks. +template +void LegacyReferenceStridedSlice( + const reference_ops::DynamicStridedSliceParams& op_params, + const RuntimeShape& input_shape, const RuntimeShape& output_shape, + SequentialTensorWriter* writer) { + ruy::profiler::ScopeLabel label("StridedSlice"); + const int dims = input_shape.DimensionsCount(); + std::vector starts(dims); + std::vector stops(dims); + std::vector input_strides(dims); + if (dims == 0) { + writer->Write(0); + return; + } + input_strides[dims - 1] = 1; + for (int i = dims - 2; i >= 0; --i) { + input_strides[i] = input_strides[i + 1] * input_shape.Dims(i + 1); + } + for (int axis = 0; axis < dims; ++axis) { + starts[axis] = reference_ops::StartForAxis(op_params, input_shape, axis); + stops[axis] = + reference_ops::EndForAxis(op_params, input_shape, axis, starts[axis]); + } + auto loop_condition = [](int64_t index, int64_t stop, int stride) { + return stride > 0 ? index < stop : index > stop; + }; + std::function write_slice = [&](int axis, + int64_t input_index) { + if (axis == dims) { + writer->Write(input_index); + return; + } + for (int64_t offset = starts[axis]; + loop_condition(offset, stops[axis], op_params.strides[axis]); + offset += op_params.strides[axis]) { + write_slice(axis + 1, input_index + offset * input_strides[axis]); + } + }; + write_slice(/*axis=*/0, /*input_index=*/0); +} + +// Dual-sentinel bit-for-bit and write-cardinality verification helper. +// 1. Runs both LegacyReferenceStridedSlice and reference_ops::StridedSlice +// twice using bitwise-complementary sentinel byte patterns (0xA5 and 0x5A). +// 2. Independently computes the closed-form analytical element count K from +// StartForAxis / EndForAxis and verifies that: +// - Every byte in [0, K * sizeof(T)) is identical between the 0xA5 run and +// the 0x5A run (proving every byte in [0, K * sizeof(T)) was written), +// - Every byte in [K * sizeof(T), max_out * sizeof(T)) equals 0xA5 in the +// first run and 0x5A in the second run (proving zero over-writes), and +// - The full output buffer matches LegacyReferenceStridedSlice bit-for-bit. +template +void VerifyBitForBitIdentical( + const reference_ops::DynamicStridedSliceParams& params, + const RuntimeShape& shape, const std::vector& input) { + const int dims = shape.DimensionsCount(); + int64_t expected_elements = 1; + for (int axis = 0; axis < dims; ++axis) { + const int64_t start = reference_ops::StartForAxis(params, shape, axis); + const int64_t stop = reference_ops::EndForAxis(params, shape, axis, start); + const int64_t stride = params.strides[axis]; + if ((stride > 0 && start >= stop) || (stride < 0 && start <= stop)) { + expected_elements = 0; + break; + } + const int64_t span = stride > 0 ? (stop - start) : (start - stop); + const int64_t step = stride > 0 ? stride : -stride; + expected_elements *= (span + step - 1) / step; + } + + const int64_t max_out = + std::max(4, static_cast(input.size()) * 2 + 4); + const size_t total_bytes = static_cast(max_out) * sizeof(T); + const size_t expected_bytes = + static_cast(expected_elements) * sizeof(T); + + std::vector opt_out_a5(max_out); + std::vector opt_out_5a(max_out); + + for (int sentinel : {0xA5, 0x5A}) { + std::vector legacy_out(max_out); + std::vector& opt_out = (sentinel == 0xA5) ? opt_out_a5 : opt_out_5a; + std::memset(legacy_out.data(), sentinel, total_bytes); + std::memset(opt_out.data(), sentinel, total_bytes); + + SequentialTensorWriter legacy_writer(input.data(), legacy_out.data()); + SequentialTensorWriter opt_writer(input.data(), opt_out.data()); + + LegacyReferenceStridedSlice(params, shape, shape, &legacy_writer); + reference_ops::StridedSlice(params, shape, shape, &opt_writer); + + EXPECT_EQ(std::memcmp(legacy_out.data(), opt_out.data(), total_bytes), 0) + << "Bit-for-bit mismatch for shape with dims=" << dims + << " and sentinel=0x" << std::hex << sentinel; + } + + if (expected_bytes > 0) { + EXPECT_EQ(std::memcmp(opt_out_a5.data(), opt_out_5a.data(), expected_bytes), + 0) + << "Under-write detected inside [0, " << expected_bytes << ")"; + } + const auto* bytes_a5 = reinterpret_cast(opt_out_a5.data()); + const auto* bytes_5a = reinterpret_cast(opt_out_5a.data()); + for (size_t b = expected_bytes; b < total_bytes; ++b) { + ASSERT_EQ(bytes_a5[b], 0xA5u) << "Over-write at byte " << b; + ASSERT_EQ(bytes_5a[b], 0x5Au) << "Over-write at byte " << b; + } +} + +template +void RunExhaustiveBitForBitSuite() { + // 0D scalar tensor. + { + RuntimeShape shape0d; + std::vector in0d = CastVector({42}); + reference_ops::DynamicStridedSliceParams params0d; + VerifyBitForBitIdentical(params0d, shape0d, in0d); + } + + // 1D, 2D, 3D, 4D, and 5D shapes exercising: + // - Zero-extent dimensions (Dims(axis) == 0) with positive & negative strides + // - Full trailing contiguous dimensions (1, 2, 3, and 4 coalesced axes) + // - Trailing singleton dimensions (Dims(axis) == 1) where + // input_strides[last_loop_axis] == 1 on a non-innermost axis + // - Partial unit-stride inner dimension (starts > 0, stops < Dims) + // - Strided inner dimension (stride = 2, 3, -1, -2, -3) + // - Strided or negative-strided outer/middle dimension with full contiguous + // inner dimensions (where input_strides[last_loop_axis] > 1) + // - Empty slices on outer, middle, and inner axes (positive & negative + // stride) + // - begin_mask, end_mask, shrink_axis_mask, and offset = true. + struct SliceCase { + std::vector dims; + std::vector starts; + std::vector stops; + std::vector strides; + uint32_t begin_mask = 0; + uint32_t end_mask = 0; + uint32_t shrink_axis_mask = 0; + bool offset = false; + }; + + const std::vector cases = { + // Zero-extent dimension cases (Dims(axis) == 0). + {{0}, {0}, {0}, {1}}, + {{0}, {0}, {0}, {-1}}, + {{0, 8, 16}, {0, 0, 0}, {0, 8, 16}, {1, 1, 1}}, + {{3, 0, 16}, {0, 0, 0}, {3, 0, 16}, {1, 1, 1}}, + {{3, 8, 0}, {0, 0, 0}, {3, 8, 0}, {1, -1, -1}}, + {{3, 0, 16}, {0, 0, 0}, {3, 0, 16}, {1, 1, 1}, 0, 0, 0b010}, + // 1D cases: full, partial, strided, reversed, empty. + {{16}, {0}, {16}, {1}}, + {{16}, {3}, {13}, {1}}, + {{16}, {1}, {15}, {2}}, + {{16}, {15}, {-1}, {-1}, 0, 1}, + {{16}, {14}, {2}, {-3}}, + {{16}, {10}, {5}, {1}}, + {{16}, {5}, {10}, {-1}}, + // 2D cases: coalesced inner, partial inner, strided outer + full inner, + // reversed outer + full inner, reversed inner. + {{6, 12}, {0, 0}, {6, 12}, {1, 1}}, + {{6, 12}, {1, 0}, {5, 12}, {1, 1}}, + {{6, 12}, {1, 2}, {5, 10}, {1, 1}}, + {{6, 12}, {0, 0}, {6, 12}, {2, 1}}, + {{6, 12}, {5, 0}, {-1, 12}, {-1, 1}, 0, 0b01}, + {{6, 12}, {5, 0}, {0, 12}, {-2, 1}}, + {{6, 12}, {1, 11}, {5, -1}, {2, -1}, 0, 0b10}, + {{6, 12}, {1, 10}, {5, 1}, {1, -2}}, + {{6, 12}, {4, 0}, {2, 12}, {1, 1}}, + {{6, 12}, {0, 8}, {6, 3}, {1, 1}}, + // Trailing singleton dimension cases (1 < inner_contig_axis < dims with + // input_strides[last_loop_axis] == 1 and last_loop_stride_is_1 == false). + {{3, 6, 1}, {0, 0, 0}, {3, 6, 1}, {1, 2, 1}}, + {{3, 6, 1}, {0, 5, 0}, {3, -1, 1}, {1, -2, 1}, 0, 0b010}, + {{2, 6, 1, 1}, {0, 5, 0, 0}, {2, -1, 1, 1}, {1, -1, 1, 1}, 0, 0b0010}, + // 3D speech-like feature tensor cases: + // [B, T, F] with full T & F coalesced, sliced T with full F coalesced, + // strided/reversed T with full F coalesced, partial F, strided F, + // and empty middle/inner axes. + {{3, 8, 16}, {0, 0, 0}, {3, 8, 16}, {1, 1, 1}}, + {{3, 8, 16}, {1, 0, 0}, {3, 8, 16}, {1, 1, 1}}, + {{3, 8, 16}, {0, 2, 0}, {3, 7, 16}, {1, 1, 1}}, + {{3, 8, 16}, {0, 1, 0}, {3, 8, 16}, {1, 2, 1}}, + {{3, 8, 16}, {0, 7, 0}, {3, -1, 16}, {1, -1, 1}, 0, 0b010}, + {{3, 8, 16}, {2, 7, 0}, {-1, 0, 16}, {-1, -2, 1}, 0, 0b001}, + {{3, 8, 16}, {0, 1, 3}, {3, 6, 14}, {1, 2, 1}}, + {{3, 8, 16}, {0, 1, 15}, {3, 6, 0}, {1, 2, -3}}, + {{3, 8, 16}, {0, 5, 0}, {3, 2, 16}, {1, 1, 1}}, + {{3, 8, 16}, {0, 1, 10}, {3, 6, 4}, {1, 1, 1}}, + {{3, 8, 16}, {1, 2, 0}, {2, 5, 16}, {1, 1, 1}, 0, 0, 0b001}, + {{3, 8, 16}, {0, 3, 0}, {3, 4, 16}, {1, 1, 1}, 0, 0, 0b010}, + {{3, 8, 16}, {0, 2, 5}, {3, 5, 6}, {1, 1, 1}, 0, 0, 0b100}, + {{3, 8, 16}, {1, 2, 3}, {2, 4, 10}, {1, 1, 1}, 0, 0, 0, true}, + // 4D and 5D cases: multi-axis coalescing across 3 and 4 trailing axes, + // mixed positive/negative strides, and shrink_axis_mask combinations. + {{2, 3, 4, 5}, {0, 1, 0, 0}, {2, 3, 4, 5}, {1, 1, 1, 1}}, + {{2, 3, 4, 5}, {0, 2, 0, 0}, {2, -1, 4, 5}, {1, -1, 1, 1}, 0, 0b0010}, + {{2, 3, 4, 5}, {1, 0, 1, 0}, {-1, 3, 3, 5}, {-1, 2, 1, 1}, 0, 0b0001}, + {{2, 2, 3, 2, 4}, {0, 0, 1, 0, 0}, {2, 2, 3, 2, 4}, {1, 1, 1, 1, 1}}, + {{2, 2, 3, 2, 4}, + {1, 1, 2, 0, 0}, + {-1, -1, 0, 2, 4}, + {-1, -1, -1, 1, 1}, + 0, + 0b00011}, + }; + + for (const auto& tc : cases) { + RuntimeShape shape(static_cast(tc.dims.size()), tc.dims.data()); + const int total_elements = shape.FlatSize(); + std::vector raw_vals(total_elements); + for (int i = 0; i < total_elements; ++i) { + if (i == 0) { + raw_vals[i] = -91; // 0xA5 in two's complement int8_t + } else if (i == 1) { + raw_vals[i] = 90; // 0x5A in int8_t + } else { + raw_vals[i] = (i * 37 + 13) % 127 - 63; + } + } + const std::vector input = CastVector(raw_vals); + reference_ops::DynamicStridedSliceParams params; + params.start_indices = tc.starts; + params.stop_indices = tc.stops; + params.strides = tc.strides; + params.begin_mask = tc.begin_mask; + params.end_mask = tc.end_mask; + params.shrink_axis_mask = tc.shrink_axis_mask; + params.offset = tc.offset; + + VerifyBitForBitIdentical(params, shape, input); + } +} + +TYPED_TEST(StridedSliceOpTest, ExhaustiveBitForBitEquivalenceWithLegacyOracle) { + RunExhaustiveBitForBitSuite(); +} + +TEST(StridedSliceBitExactInt64Test, + ExhaustiveBitForBitEquivalenceWithLegacyOracleInt64) { + RunExhaustiveBitForBitSuite(); +} + +// Verify strict bit-for-bit equality on float tensors containing special IEEE +// 754 bit patterns (-0.0f, +0.0f, quiet NaNs with distinct payloads, signaling +// NaNs, subnormals, +-Inf, and sentinel bit patterns 0xa5a5a5a5 / 0x5a5a5a5a) +// where operator== would either return false (NaN != NaN) or conflate distinct +// bit representations (-0.0f == +0.0f). +TEST(StridedSliceBitExactFloatTest, PreservesSpecialFloatBitPatternsExactly) { + const std::vector special_bits = { + 0x00000000u, // +0.0f + 0x80000000u, // -0.0f + 0x7f800000u, // +Inf + 0xff800000u, // -Inf + 0x7fc00001u, // Quiet NaN payload 1 + 0x7fcabcdeu, // Quiet NaN payload 0xabcde (8 hex digits) + 0xffc01234u, // Negative Quiet NaN payload 0x1234 + 0x7f800001u, // Signaling NaN payload 1 + 0x00000001u, // Smallest positive subnormal + 0x80000001u, // Smallest negative subnormal + 0x007fffffu, // Largest positive subnormal + 0x3f800000u, // 1.0f + 0xbf800000u, // -1.0f + 0xa5a5a5a5u, // Sentinel bit pattern A5 + 0x5a5a5a5au, // Sentinel bit pattern 5A + }; + + RuntimeShape shape({2, 4, 8}); + const int total = shape.FlatSize(); + std::vector input(total); + for (int i = 0; i < total; ++i) { + uint32_t bits = + special_bits[i % special_bits.size()] ^ static_cast(i & 0x3); + std::memcpy(&input[i], &bits, sizeof(float)); + } + + for (int stride_1 : {1, 2, -1, -2}) { + for (int stride_2 : {1, 2, -1}) { + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, stride_1 > 0 ? 0 : 3, stride_2 > 0 ? 0 : 7}; + params.stop_indices = {2, stride_1 > 0 ? 4 : -1, stride_2 > 0 ? 8 : -1}; + params.strides = {1, stride_1, stride_2}; + params.end_mask = + (stride_1 < 0 ? 0b010u : 0u) | (stride_2 < 0 ? 0b100u : 0u); + VerifyBitForBitIdentical(params, shape, input); + } + } +} + +// Verify bit-for-bit identical serialized string buffers on multi-dimensional +// std::string tensors across coalesced, partial, strided/reversed, and empty +// slices, comparing both deserialized strings and raw serialized TfLiteTensor +// byte buffers against LegacyReferenceStridedSlice. +TEST(StridedSliceBitExactStringTest, MultiDimCoalescedAndStridedStrings) { + std::vector input_data; + input_data.reserve(24); + for (int i = 0; i < 24; ++i) { + if (i % 5 == 0) { + input_data.push_back(""); // Empty string + } else if (i % 7 == 0) { + input_data.push_back(std::string("nul\0byte_", 9) + std::to_string(i)); + } else { + input_data.push_back("token_str_" + std::to_string(i * 101)); + } + } + + auto verify_raw_string_tensor_bit_exact = + [&](const RuntimeShape& shape, + const reference_ops::DynamicStridedSliceParams& params) { + DynamicBuffer in_buf; + for (const std::string& s : input_data) { + in_buf.AddString(s.data(), s.size()); + } + TfLiteTensor input_tensor{}; + input_tensor.type = kTfLiteString; + input_tensor.allocation_type = kTfLiteDynamic; + in_buf.WriteToTensor(&input_tensor, /*new_shape=*/nullptr); + + TfLiteTensor legacy_tensor{}; + legacy_tensor.type = kTfLiteString; + legacy_tensor.allocation_type = kTfLiteDynamic; + + TfLiteTensor opt_tensor{}; + opt_tensor.type = kTfLiteString; + opt_tensor.allocation_type = kTfLiteDynamic; + + { + SequentialTensorWriter legacy_writer(&input_tensor, + &legacy_tensor); + LegacyReferenceStridedSlice(params, shape, shape, + &legacy_writer); + } + { + SequentialTensorWriter opt_writer(&input_tensor, + &opt_tensor); + reference_ops::StridedSlice(params, shape, shape, + &opt_writer); + } + + ASSERT_EQ(legacy_tensor.bytes, opt_tensor.bytes); + if (legacy_tensor.bytes > 0) { + EXPECT_EQ(std::memcmp(legacy_tensor.data.raw, opt_tensor.data.raw, + legacy_tensor.bytes), + 0); + } + + free(input_tensor.data.raw); + free(legacy_tensor.data.raw); + free(opt_tensor.data.raw); + + // Also verify a non-trivially-copyable type with the raw-pointer + // overload (SequentialTensorWriter(const T*, T*)), which exercises + // the !std::is_trivially_copyable_v element-by-element Write + // path in write_contiguous without invoking memcpy. + struct NonTriviallyCopyableVal { + int v = -1; + const NonTriviallyCopyableVal* self = this; + NonTriviallyCopyableVal() = default; + explicit NonTriviallyCopyableVal(int x) : v(x), self(this) {} + NonTriviallyCopyableVal(const NonTriviallyCopyableVal& o) + : v(o.v), self(this) {} + NonTriviallyCopyableVal& operator=(const NonTriviallyCopyableVal& o) { + v = o.v; + self = this; + return *this; + } + bool operator==(const NonTriviallyCopyableVal& o) const { + return v == o.v && self == this && o.self == &o; + } + }; + static_assert(!std::is_trivially_copyable_v, + "Must be non-trivially copyable"); + std::vector ntc_in(shape.FlatSize()); + for (int i = 0; i < shape.FlatSize(); ++i) { + ntc_in[i] = NonTriviallyCopyableVal(i + 100); + } + std::vector ntc_legacy( + shape.FlatSize(), NonTriviallyCopyableVal(-999)); + std::vector ntc_actual( + shape.FlatSize(), NonTriviallyCopyableVal(-999)); + SequentialTensorWriter ntc_legacy_writer( + ntc_in.data(), ntc_legacy.data()); + LegacyReferenceStridedSlice( + params, shape, shape, &ntc_legacy_writer); + reference_ops::StridedSlice( + params, shape, ntc_in.data(), shape, ntc_actual.data()); + EXPECT_EQ(ntc_actual, ntc_legacy); + }; + + const RuntimeShape shape3d({2, 3, 4}); + + // Case 1: Coalesced inner dimension [2, 3, 4] -> [0:2, 1:3, 0:4] + { + StridedSliceOpModel m({2, 3, 4}, {3}, {3}, {3}, input_data, + {0, 1, 0}, {2, 3, 4}, {1, 1, 1}, 0, 0, 0, + 0, 0, false); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 4})); + std::vector expected; + for (int b = 0; b < 2; ++b) { + for (int r = 1; r < 3; ++r) { + for (int c = 0; c < 4; ++c) { + expected.push_back(input_data[b * 12 + r * 4 + c]); + } + } + } + EXPECT_THAT(m.GetStringOutput(), ElementsAreArray(expected)); + + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, 1, 0}; + params.stop_indices = {2, 3, 4}; + params.strides = {1, 1, 1}; + verify_raw_string_tensor_bit_exact(shape3d, params); + } + + // Case 2: Negative stride on middle axis with full unit-stride inner axis + // (exercises input_strides[last_loop_axis] > 1 with WriteN on strings). + { + StridedSliceOpModel m({2, 3, 4}, {3}, {3}, {3}, input_data, + {0, 2, 0}, {2, -1, 4}, {1, -1, 1}, 0, + 0b010, 0, 0, 0, false); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3, 4})); + std::vector expected; + for (int b = 0; b < 2; ++b) { + for (int r = 2; r >= 0; --r) { + for (int c = 0; c < 4; ++c) { + expected.push_back(input_data[b * 12 + r * 4 + c]); + } + } + } + EXPECT_THAT(m.GetStringOutput(), ElementsAreArray(expected)); + + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, 2, 0}; + params.stop_indices = {2, -1, 4}; + params.strides = {1, -1, 1}; + params.end_mask = 0b010; + verify_raw_string_tensor_bit_exact(shape3d, params); + } + + // Case 3: Partial unit-stride inner axis (1:3:1), strided/reversed inner axis + // (3:-1:-1), and empty string slice (2:1:1). + { + reference_ops::DynamicStridedSliceParams partial_inner; + partial_inner.start_indices = {0, 0, 1}; + partial_inner.stop_indices = {2, 3, 3}; + partial_inner.strides = {1, 1, 1}; + verify_raw_string_tensor_bit_exact(shape3d, partial_inner); + + reference_ops::DynamicStridedSliceParams reversed_inner; + reversed_inner.start_indices = {0, 0, 3}; + reversed_inner.stop_indices = {2, 3, -1}; + reversed_inner.strides = {1, 2, -1}; + reversed_inner.end_mask = 0b100; + verify_raw_string_tensor_bit_exact(shape3d, reversed_inner); + + reference_ops::DynamicStridedSliceParams empty_slice; + empty_slice.start_indices = {0, 2, 0}; + empty_slice.stop_indices = {2, 1, 4}; + empty_slice.strides = {1, 1, 1}; + verify_raw_string_tensor_bit_exact(shape3d, empty_slice); + } +} + +// Microbenchmarks comparing LegacyReferenceStridedSlice vs optimized +// reference_ops::StridedSlice on representative speech_detector_alt_service +// feature tensors ([1, 100, 256] and [4, 64, 128]). +void BM_StridedSlice_SpeechFeature_Legacy(benchmark::State& state) { + const RuntimeShape input_shape({1, 100, 256}); + const RuntimeShape output_shape({1, 99, 256}); + std::vector input(input_shape.FlatSize(), 1.25f); + std::vector output(output_shape.FlatSize(), 0.0f); + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, 1, 0}; + params.stop_indices = {1, 100, 256}; + params.strides = {1, 1, 1}; + + for (auto _ : state) { + SequentialTensorWriter writer(input.data(), output.data()); + LegacyReferenceStridedSlice(params, input_shape, output_shape, + &writer); + benchmark::DoNotOptimize(output.data()); + benchmark::ClobberMemory(); + } +} +BENCHMARK(BM_StridedSlice_SpeechFeature_Legacy)->MinTime(0.01); + +void BM_StridedSlice_SpeechFeature_Optimized(benchmark::State& state) { + const RuntimeShape input_shape({1, 100, 256}); + const RuntimeShape output_shape({1, 99, 256}); + std::vector input(input_shape.FlatSize(), 1.25f); + std::vector output(output_shape.FlatSize(), 0.0f); + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, 1, 0}; + params.stop_indices = {1, 100, 256}; + params.strides = {1, 1, 1}; + + for (auto _ : state) { + SequentialTensorWriter writer(input.data(), output.data()); + reference_ops::StridedSlice(params, input_shape, output_shape, + &writer); + benchmark::DoNotOptimize(output.data()); + benchmark::ClobberMemory(); + } +} +BENCHMARK(BM_StridedSlice_SpeechFeature_Optimized)->MinTime(0.01); + +void BM_StridedSlice_BatchedSpeechFeature_Legacy(benchmark::State& state) { + const RuntimeShape input_shape({4, 64, 128}); + const RuntimeShape output_shape({4, 60, 128}); + std::vector input(input_shape.FlatSize(), 1.25f); + std::vector output(output_shape.FlatSize(), 0.0f); + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, 2, 0}; + params.stop_indices = {4, 62, 128}; + params.strides = {1, 1, 1}; + + for (auto _ : state) { + SequentialTensorWriter writer(input.data(), output.data()); + LegacyReferenceStridedSlice(params, input_shape, output_shape, + &writer); + benchmark::DoNotOptimize(output.data()); + benchmark::ClobberMemory(); + } +} +BENCHMARK(BM_StridedSlice_BatchedSpeechFeature_Legacy)->MinTime(0.01); + +void BM_StridedSlice_BatchedSpeechFeature_Optimized(benchmark::State& state) { + const RuntimeShape input_shape({4, 64, 128}); + const RuntimeShape output_shape({4, 60, 128}); + std::vector input(input_shape.FlatSize(), 1.25f); + std::vector output(output_shape.FlatSize(), 0.0f); + reference_ops::DynamicStridedSliceParams params; + params.start_indices = {0, 2, 0}; + params.stop_indices = {4, 62, 128}; + params.strides = {1, 1, 1}; + + for (auto _ : state) { + SequentialTensorWriter writer(input.data(), output.data()); + reference_ops::StridedSlice(params, input_shape, output_shape, + &writer); + benchmark::DoNotOptimize(output.data()); + benchmark::ClobberMemory(); + } +} +BENCHMARK(BM_StridedSlice_BatchedSpeechFeature_Optimized)->MinTime(0.01); + } // namespace } // namespace tflite From 89ba87d707f0006ae9da367a31366120ed2e0f65 Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Mon, 28 Sep 2026 01:27:47 -0700 Subject: [PATCH 21/25] PR #49586: [ROCm]Add rocm_platform_test. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/49586 ๐Ÿ“ Summary of Changes - Add `//xla/stream_executor/rocm:rocm_platform_test`, the ROCm twin of `cuda_platform_test` / `sycl_platform_test`. - Cover `PlatformWithName("ROCM")` and `FindExisting` after `ExecutorForDevice`. ๐ŸŽฏ Justification Add more unit test coverage on ROCm platform. ๐Ÿš€ Kind of Contribution ๐Ÿงช Tests Copybara import of the project: -- e712cca80234d5417a06533eee3fff9fa9fdd633 by linchen1 : Add rocm_platform_test. Signed-off-by: linchen1 Merging this change closes #49586 PiperOrigin-RevId: 989479036 --- .../xla/xla/stream_executor/rocm/BUILD | 20 +++++++ .../rocm/rocm_platform_test.cc | 57 +++++++++++++++++++ 2 files changed, 77 insertions(+) create mode 100644 third_party/xla/xla/stream_executor/rocm/rocm_platform_test.cc diff --git a/third_party/xla/xla/stream_executor/rocm/BUILD b/third_party/xla/xla/stream_executor/rocm/BUILD index f95ca29577e1c5..2632342565b810 100644 --- a/third_party/xla/xla/stream_executor/rocm/BUILD +++ b/third_party/xla/xla/stream_executor/rocm/BUILD @@ -333,6 +333,26 @@ cc_library( alwayslink = True, # Registers itself with the PlatformManager. ) +xla_test( + name = "rocm_platform_test", + srcs = ["rocm_platform_test.cc"], + backends = ["gpu"], + tags = [ + "gpu", + "rocm-only", + ], + deps = [ + ":rocm_platform", # buildcleaner: keep + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream_executor_h", + "//xla/tsl/platform:statusor", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", # buildcleaner: keep + ], +) + cc_library( name = "rocm_runtime_abi_version", srcs = ["rocm_runtime_abi_version.cc"], diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_platform_test.cc b/third_party/xla/xla/stream_executor/rocm/rocm_platform_test.cc new file mode 100644 index 00000000000000..5e1a490151a84a --- /dev/null +++ b/third_party/xla/xla/stream_executor/rocm/rocm_platform_test.cc @@ -0,0 +1,57 @@ +/* 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/stream_executor/rocm/rocm_platform.h" + +#include +#include +#include "absl/container/flat_hash_map.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tsl/platform/statusor.h" + +namespace stream_executor::gpu { +namespace { + +TEST(RocmPlatformTest, TestPlatformName) { + ASSERT_OK_AND_ASSIGN(Platform * platform, + PlatformManager::PlatformWithName("ROCM")); + EXPECT_EQ(platform->Name(), "ROCM"); +} + +// Miss before ExecutorForDevice; same pointer after. +TEST(RocmPlatformTest, FindExistingWorks) { + ASSERT_OK_AND_ASSIGN(Platform * platform, + PlatformManager::PlatformWithName("ROCM")); + const int num_devices = platform->VisibleDeviceCount(); + ASSERT_GT(num_devices, 0); + for (int i = 0; i < num_devices; ++i) { + EXPECT_FALSE(platform->FindExisting(i).ok()); + } + absl::flat_hash_map executors; + for (int i = 0; i < num_devices; ++i) { + ASSERT_OK_AND_ASSIGN(auto executor, platform->ExecutorForDevice(i)); + executors[i] = executor; + } + EXPECT_EQ(executors.size(), num_devices); + for (int i = 0; i < num_devices; ++i) { + ASSERT_OK_AND_ASSIGN(auto executor, platform->FindExisting(i)); + EXPECT_EQ(executor, executors[i]); + } +} + +} // namespace +} // namespace stream_executor::gpu From 9725c54ec837e8994c152e0e45d8b699c52ec55a Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Mon, 28 Sep 2026 01:32:42 -0700 Subject: [PATCH 22/25] [XLA:GPU] Fix out-of-bounds index handling in ScatterDeterminismExpander Previously, `ScatterDeterminismExpander` only checked for out-of-bounds scatter indices when updates were non-scalar. For scalar updates with multi-dimensional scatter indices, `SortIndicesAndUpdates` flattened the indices via a dot product with the operand strides without validating each component. Out-of-bounds coordinates in one dimension (e.g., `(0, 4)` or `(1, -4)` on a `[2, 4]` operand) could wrap into adjacent rows and overwrite valid elements instead of being dropped. Additionally, for non-scalar updates with implicit dimensions, `out_of_bound_tensor` was padded with `0`s for implicit dimensions via `AddImplicitDimensionsToIndices`. When `ExpandIndices` added non-zero window offsets along an implicit dimension, the flattened out-of-bounds index could collide with valid in-bounds indices. Validate multi-dimensional scatter indices and mask out-of-bounds rows with `out_of_bound_tensor` in the scalar-update path before sorting, and recreate `out_of_bound_tensor` via `CreateBoundTensor` using the full index-to-operand dimension mapping when implicit dimensions are added. Fixes #49380 PiperOrigin-RevId: 989480847 --- .../scatter_determinism_expander.cc | 27 ++++-- .../scatter_determinism_expander_test.cc | 84 +++++++++++++++++++ 2 files changed, 105 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.cc b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.cc index f48ed029f55ef1..884adaa18bb3e6 100644 --- a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.cc +++ b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander.cc @@ -788,12 +788,14 @@ absl::StatusOr ScatterDeterminismExpander::ExpandInstruction( dim_numbers.scatter_dims_to_operand_dims(), scatter_indices)); CHECK(scatter_indices->shape().dimensions(0) == scatter_indices_count); - // Add implicit dimensions to OOB constant, if needed. - ABSL_ASSIGN_OR_RETURN( - out_of_bound_tensor, - AddImplicitDimensionsToIndices( - scatter_operands[0]->shape().dimensions().size(), - dim_numbers.scatter_dims_to_operand_dims(), out_of_bound_tensor)); + // Recompute OOB constant with implicit dimensions, if needed. + if (operand_shape.dimensions().size() != + dim_numbers.scatter_dims_to_operand_dims_size()) { + ABSL_ASSIGN_OR_RETURN( + out_of_bound_tensor, + CreateBoundTensor(parent, scatter_indices, operand_shape.dimensions(), + full_index_to_operand_dims)); + } // If any updates are out of bound, we change the corresponding indices to // be oob_tensor values @@ -841,6 +843,19 @@ absl::StatusOr ScatterDeterminismExpander::ExpandInstruction( full_index_to_operand_dims[i]); } } else { + if (!has_scalar_indices) { + std::vector actual_update_window_dims( + scatter_operands[0]->shape().dimensions().size(), 1); + ABSL_ASSIGN_OR_RETURN( + HloInstruction * oob_check_mask, + CheckValidIndices(parent, scatter_indices, + scatter_operands[0]->shape().dimensions(), + actual_update_window_dims, + dim_numbers.scatter_dims_to_operand_dims())); + scatter_indices = parent->AddInstruction(HloInstruction::CreateTernary( + scatter_indices->shape(), HloOpcode::kSelect, oob_check_mask, + scatter_indices, out_of_bound_tensor)); + } new_dim_numbers = dim_numbers; } diff --git a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc index 935ceb79c7efe5..0bd48325f49063 100644 --- a/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/scatter_determinism_expander_test.cc @@ -458,6 +458,90 @@ TEST_P(ScatterDeterminismExpanderTest, EXPECT_EQ(actual_result, expected_result); } +TEST_P(ScatterDeterminismExpanderTest, + ScatterAddWithNonScalarIndexOutOfBoundCorrectnessTest) { + auto index_type = GetParam(); + const char* const kModuleTemplate = R"( + HloModule scatter_determinism_expander + + scatter_computation { + arg1.173 = f32[] parameter(1) + arg0.172 = f32[] parameter(0) + ROOT add.48 = f32[] add(arg0.172, arg1.173) + } + + ENTRY scatter_add_computation { + operand = f32[2, 4] constant({{0, 0, 0, 0}, {0, 0, 0, 0}}) + indices = $0[5, 2] constant({{0, 0}, {0, 4}, {1, 0}, {1, 4}, {1, -4}}) + updates = f32[5] constant({1, 2, 3, 4, 5}) + ROOT scatter.48 = f32[2, 4] scatter(operand, indices, updates), + update_window_dims={}, inserted_window_dims={0, 1}, + scatter_dims_to_operand_dims={0, 1}, index_vector_dim=1, + to_apply=scatter_computation + })"; + + const std::string hlo = absl::Substitute( + kModuleTemplate, primitive_util::LowercasePrimitiveTypeName(index_type)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + auto cloned_module = module->Clone(); + ASSERT_OK_AND_ASSIGN(Literal expected_literal, + Execute(std::move(cloned_module), {})); + auto expected_result = expected_literal.data(); + + ASSERT_OK_AND_ASSIGN(bool result, + RunScatterDeterminismExpander(module.get())); + EXPECT_TRUE(result); + + ASSERT_OK_AND_ASSIGN(Literal result_literal, Execute(std::move(module), {})); + + auto result_data = result_literal.data(); + std::vector actual_result(result_data.begin(), result_data.end()); + + EXPECT_EQ(actual_result, expected_result); +} + +TEST_P(ScatterDeterminismExpanderTest, + ScatterAddWithNonScalarUpdateAndImplicitDimensionsOutOfBoundTest) { + auto index_type = GetParam(); + const char* const kModuleTemplate = R"( + HloModule scatter_determinism_expander + + scatter_computation { + arg1.173 = f32[] parameter(1) + arg0.172 = f32[] parameter(0) + ROOT add.48 = f32[] add(arg0.172, arg1.173) + } + + ENTRY scatter_add_computation { + operand = f32[2, 4] constant({{0, 0, 0, 0}, {0, 0, 0, 0}}) + indices = $0[3] constant({0, 4, -1}) + updates = f32[3, 2] constant({{1, 2}, {10, 20}, {100, 200}}) + ROOT scatter.48 = f32[2, 4] scatter(operand, indices, updates), + update_window_dims={1}, inserted_window_dims={1}, + scatter_dims_to_operand_dims={1}, index_vector_dim=1, + to_apply=scatter_computation + })"; + + const std::string hlo = absl::Substitute( + kModuleTemplate, primitive_util::LowercasePrimitiveTypeName(index_type)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + auto cloned_module = module->Clone(); + ASSERT_OK_AND_ASSIGN(Literal expected_literal, + Execute(std::move(cloned_module), {})); + auto expected_result = expected_literal.data(); + + ASSERT_OK_AND_ASSIGN(bool result, + RunScatterDeterminismExpander(module.get())); + EXPECT_TRUE(result); + + ASSERT_OK_AND_ASSIGN(Literal result_literal, Execute(std::move(module), {})); + + auto result_data = result_literal.data(); + std::vector actual_result(result_data.begin(), result_data.end()); + + EXPECT_EQ(actual_result, expected_result); +} + TEST_P(ScatterDeterminismExpanderTest, ScatterAddWithNonScalarIndexAndUpdateCorrectness2DTest1) { auto index_type = GetParam(); From e1a22a956a2f8360c9b9345f4df39b8be4582278 Mon Sep 17 00:00:00 2001 From: Dirk Hornung Date: Mon, 28 Sep 2026 01:33:45 -0700 Subject: [PATCH 23/25] Mutex guard cuDNN compilation handle. The autotuner shares the cuDNN handle by multiple threads potentially causing race conditions. From the documentation (https://docs.nvidia.com/deeplearning/cudnn/latest/developer/misc.html#thread-safety): "The cuDNN library is thread-safe. Its functions can be called from multiple host threads, so long as the threads do not share the same cuDNN handle simultaneously." PiperOrigin-RevId: 989481298 --- .../xla/xla/stream_executor/cuda/cuda_dnn.cc | 41 +++++++++++++------ 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc index b45a2271673e4b..b4321feb17adfe 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc @@ -202,7 +202,8 @@ std::string CudnnStatusToString(cudnnStatus_t status) { // RAII wrapper for all calls to cuDNN with a cuDNN handle argument. // -// See CudnnAccess::GetHandle() for details. +// See CudnnAccess::GetHandle() and CudnnAccess::GetCompilationHandle() for +// details. class CudnnHandle { public: // Takes ownership of the lock to access cuDNN using handle. @@ -212,6 +213,11 @@ class CudnnHandle { lock_(std::move(lock)), handle_(handle) {} + // Takes ownership of the lock to access cuDNN using handle. Doesn't activate + // a CUDA context. + CudnnHandle(std::unique_ptr lock, cudnnHandle_t handle) + : lock_(std::move(lock)), handle_(handle) {} + // Returns cuDNN handle. To be passed directly to cuDNN APIs, don't keep // a copy. cudnnHandle_t handle() const { return handle_; } @@ -249,6 +255,7 @@ class CudnnAccess { explicit CudnnAccess(cudnnHandle_t handle) : handle_(handle) {} absl::Status InitializeCompilationHandle() { + absl::MutexLock lock(compilation_mutex_); if (cudnnCreate(&compilation_handle_) != CUDNN_STATUS_SUCCESS) { return absl::InternalError( "Creation of compilation cudnn handle failed."); @@ -257,7 +264,6 @@ class CudnnAccess { } ~CudnnAccess() { - absl::MutexLock lock(mutex_); cudnnDestroy(handle_); if (compilation_handle_) { @@ -295,11 +301,15 @@ class CudnnAccess { return CudnnHandle(executor, std::move(lock), handle_); } - absl::StatusOr GetCompilationHandle() { + // Creates a CudnnHandle instance for the compilation handle, which is used + // to build cuDNN graphs and execution plans. + absl::StatusOr GetCompilationHandle() { + auto lock = std::make_unique(compilation_mutex_); + compilation_mutex_.AssertHeld(); if (!compilation_handle_) { return absl::InternalError("CudnnAccess not properly initialized."); } - return compilation_handle_; + return CudnnHandle(std::move(lock), compilation_handle_); } void NotifyStreamDestroyed(Stream* stream) { @@ -323,8 +333,13 @@ class CudnnAccess { // cuDNN library handle. cudnnHandle_t handle_ ABSL_GUARDED_BY(mutex_); // Owned. - // Shared compilation handle for all threads calling GetCompilationHandle() - cudnnHandle_t compilation_handle_ = nullptr; + // Guards the use of compilation_handle_ below. Separate from mutex_ so that + // compilation doesn't block execution. + absl::Mutex compilation_mutex_; + + // Shared compilation handle for all threads calling GetCompilationHandle(). + cudnnHandle_t compilation_handle_ ABSL_GUARDED_BY(compilation_mutex_) = + nullptr; // Owned. }; namespace { @@ -6888,12 +6903,14 @@ absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support, if (dnn_support) { const CudnnSupport& cudnn_support = static_cast(*dnn_support); - ABSL_ASSIGN_OR_RETURN(auto cudnn_handle, + // Holds the lock on the shared compilation handle until the end of scope. + ABSL_ASSIGN_OR_RETURN(CudnnHandle cudnn, cudnn_support.cudnn_->GetCompilationHandle()); RETURN_IF_CUDNN_FRONTEND_ERROR(graph_.validate()); - RETURN_IF_CUDNN_FRONTEND_ERROR(graph_.build_operation_graph(cudnn_handle)); + RETURN_IF_CUDNN_FRONTEND_ERROR( + graph_.build_operation_graph(cudnn.handle())); ABSL_RETURN_IF_ERROR(create_and_filter_plans()); - RETURN_CUDNN_FRONTEND_STATUS(graph_.check_support(cudnn_handle)); + RETURN_CUDNN_FRONTEND_STATUS(graph_.check_support(cudnn.handle())); } else { // Deviceless mode. No cuDNN version guard needed: DeviceProperties // deserialization inside BuildDeviceProperties rejects runtimes < 9.8. @@ -6914,13 +6931,13 @@ absl::Status CudnnGraph::Build(dnn::DnnSupport* dnn_support, if (dnn_support) { const CudnnSupport& cudnn_support = static_cast(*dnn_support); - ABSL_ASSIGN_OR_RETURN(auto cudnn_handle, + ABSL_ASSIGN_OR_RETURN(CudnnHandle cudnn, cudnn_support.cudnn_->GetCompilationHandle()); if (plan_id.has_value()) { RETURN_CUDNN_FRONTEND_STATUS( - graph_.build_plan_at_index(cudnn_handle, *plan_id)); + graph_.build_plan_at_index(cudnn.handle(), *plan_id)); } - RETURN_CUDNN_FRONTEND_STATUS(graph_.build_plans(cudnn_handle)); + RETURN_CUDNN_FRONTEND_STATUS(graph_.build_plans(cudnn.handle())); } else { // no need to set_device_properties, it is done in Prepare() if (plan_id.has_value()) { From 96aebc239a6331be2b094b5bff16e09ab548b26d Mon Sep 17 00:00:00 2001 From: Dirk Hornung Date: Mon, 28 Sep 2026 02:11:09 -0700 Subject: [PATCH 24/25] Add constant folding pass before fusing convolutions. CuDNN graphs do not support the following pattern ``` const = bf16[] constant(6) conv = f32[] convert(const) ``` PiperOrigin-RevId: 989496996 --- third_party/xla/xla/service/gpu/BUILD | 2 ++ third_party/xla/xla/service/gpu/fusion_pipeline.cc | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 9a1d99fca60280..e671287b8aa1ff 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -1709,6 +1709,7 @@ cc_library( hdrs = ["fusion_pipeline.h"], deps = [ ":alias_info", + "//xla:shape_util", "//xla:xla_proto_cc", "//xla/backends/gpu/transforms:conv_canonicalizer", "//xla/backends/gpu/transforms:conv_fusion_rewriter", @@ -1720,6 +1721,7 @@ cc_library( "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", "//xla/hlo/pass:hlo_pass_pipeline", + "//xla/hlo/transforms/simplifiers:hlo_constant_folding", "//xla/hlo/transforms/simplifiers:hlo_dce", "//xla/service:cpu_gpu_shape_verifier", "//xla/service:hlo_cost_analysis", diff --git a/third_party/xla/xla/service/gpu/fusion_pipeline.cc b/third_party/xla/xla/service/gpu/fusion_pipeline.cc index 5c4ce8c25203fd..e52e7f25282a1f 100644 --- a/third_party/xla/xla/service/gpu/fusion_pipeline.cc +++ b/third_party/xla/xla/service/gpu/fusion_pipeline.cc @@ -29,6 +29,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/pass/hlo_pass_fix.h" #include "xla/hlo/pass/hlo_pass_pipeline.h" +#include "xla/hlo/transforms/simplifiers/hlo_constant_folding.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/service/cpu_gpu_shape_verifier.h" #include "xla/service/gpu/alias_info.h" @@ -37,6 +38,8 @@ limitations under the License. #include "xla/service/hlo_cse.h" #include "xla/service/hlo_verifier.h" #include "xla/service/layout_assignment.h" +#include "xla/shape.h" +#include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" #include "xla/tsl/platform/threadpool.h" #include "xla/xla.pb.h" @@ -67,6 +70,12 @@ HloPassPipeline FusionPipeline( // Rewrite convs into conv fusions. if (!debug_options.xla_gpu_experimental_disable_binary_libraries() && debug_options.xla_gpu_experimental_enable_conv_fusion()) { + HloConstantFolding::Options constant_folding_options; + constant_folding_options.is_layout_sensitive = true; + constant_folding_options.can_fold_shape = [](const Shape& shape) { + return ShapeUtil::IsEffectiveScalar(shape); + }; + fusion.AddPass(constant_folding_options); fusion.AddPass(); fusion.AddPass(gpu_device_info); } From 57e62153517e7d2c911651507049334d0b5d4a9b Mon Sep 17 00:00:00 2001 From: Mohammed Anany Date: Mon, 28 Sep 2026 02:26:38 -0700 Subject: [PATCH 25/25] [XLA:GPU] Allow a symbolic TilingSpace to be cloned into another MLIRContext This allows later changes to evaluate tiling candidates concurrently, each in its own context. This is essential because mlir contexts in GpuCompiler are not thread-safe. PiperOrigin-RevId: 989503240 --- .../tiling/experimental/tiling_space.cc | 48 ++++++++++++++++--- .../tiling/experimental/tiling_space.h | 14 +++++- .../tiling/experimental/tiling_space_test.cc | 36 ++++++++++++++ 3 files changed, 90 insertions(+), 8 deletions(-) diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc index f010e9db706c48..2cfe2bb9b0521a 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc @@ -538,10 +538,8 @@ absl::StatusOr> TilingSpace::Create( for (auto [index, dim] : llvm::enumerate(dims)) { int64_t global_dim_id = tiling_space->GetDimensionInfo(root.instruction(), index).id.value(); - dim_tiles.push_back(GetDefaultDimTile( - TiledDimId(global_dim_id), - CreateSymbolExpr(global_dim_id, tiling_space->num_dimensions(), ctx), - dim)); + dim_tiles.push_back( + tiling_space->GetDefaultRootDimTile(TiledDimId(global_dim_id), dim)); } Tile tile{*tiling_space, std::move(dim_tiles)}; if (root_shape.IsTuple()) { @@ -556,9 +554,24 @@ absl::StatusOr> TilingSpace::Create( return tiling_space; } -std::unique_ptr TilingSpace::Clone() const { +std::unique_ptr TilingSpace::Clone( + mlir::MLIRContext* target_context) const { + const bool rebind = + target_context != nullptr && target_context != mlir_context_; + if (rebind) { + // Only the default root tiles of a symbolic space can be rebuilt in another + // context. Constraints hold expressions of this space's context, so they + // must be trivial. + CHECK(is_symbolic_) << "Cloning into another MLIRContext is only " + "supported for a symbolic TilingSpace."; + CHECK(divisibility_constraints_.empty() && constraint_.IsAlwaysSatisfied()) + << "Cloning a TilingSpace with constraints into another MLIRContext is " + "not supported."; + RegisterSymbolicExprStorage(target_context); + } + auto cloned = std::make_unique(); - cloned->mlir_context_ = mlir_context_; + cloned->mlir_context_ = rebind ? target_context : mlir_context_; cloned->is_symbolic_ = is_symbolic_; cloned->constraint_ = constraint_; cloned->divisibility_constraints_ = divisibility_constraints_; @@ -580,7 +593,21 @@ std::unique_ptr TilingSpace::Clone() const { cloned->tiled_roots_.reserve(tiled_roots_.size()); for (const auto& root_tile : tiled_roots_) { - cloned->tiled_roots_.push_back(root_tile.CloneWithNewTilingSpace(*cloned)); + if (!rebind) { + cloned->tiled_roots_.push_back( + root_tile.CloneWithNewTilingSpace(*cloned)); + continue; + } + // A symbolic root tile is the default one built by Create, whose size is + // the symbol of its dimension. + llvm::SmallVector dim_tiles; + dim_tiles.reserve(root_tile.dim_tiles().size()); + for (const DimTile& dim_tile : root_tile.dim_tiles()) { + dim_tiles.push_back(cloned->GetDefaultRootDimTile( + TiledDimId(dim_tile.size.GetValue() - num_dimensions()), + dim_tile.upper_bound.GetValue())); + } + cloned->tiled_roots_.push_back(Tile{*cloned, std::move(dim_tiles)}); } cloned->dim_vars_indexing_ = dim_vars_indexing_; @@ -590,6 +617,13 @@ std::unique_ptr TilingSpace::Clone() const { return cloned; } +DimTile TilingSpace::GetDefaultRootDimTile(TiledDimId id, + int64_t dim_size) const { + return GetDefaultDimTile( + id, CreateSymbolExpr(id.value(), num_dimensions(), mlir_context_), + dim_size); +} + int64_t TilingSpace::num_parallel_dimensions() const { return absl::c_count_if(dimensions_, [](const DimensionInfo& dim) { return dim.type == DimensionSemantics::kParallel; diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h index 4c1c5a24b50234..79848aad505c3c 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h @@ -135,7 +135,15 @@ class TilingSpace { // Creates an independent deep copy of the TilingSpace, with all internal // pointer maps and root tiles re-bound to the new instance. - std::unique_ptr Clone() const; + // + // If `target_context` is null or is this space's MLIRContext, the copy + // shares this space's context and symbolic expressions. Otherwise, the root + // tiles are rebuilt in `target_context`, so that the copy is fully + // independent of this space's context and can be tiled concurrently with it. + // + // REQUIRES: IsSymbolic() if `target_context` is a different context. + std::unique_ptr Clone( + mlir::MLIRContext* target_context = nullptr) const; std::string ToString() const; @@ -223,6 +231,10 @@ class TilingSpace { // building indexing maps during simplification. void InitSimplificationIndexing(); + // Returns the default symbolic tile, in this space's context, for the root + // dimension `id` of size `dim_size`. + DimTile GetDefaultRootDimTile(TiledDimId id, int64_t dim_size) const; + // Maps from (hlo, dim_position) to the dimension info. absl::flat_hash_map, const DimensionInfo*> diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc index 8bb35fcda03251..45a05967886d03 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc @@ -645,5 +645,41 @@ TEST_F(TilingSpaceTest, CloneAssignsIndependentTileSizes) { EXPECT_EQ(cloned_space->dimensions()[1].tile_size, 2); } +TEST_F(TilingSpaceTest, CloneIntoAnotherContextRebindsRootTiles) { + auto root = ParseAndGetRoot(R"( + HloModule m + ENTRY e { + p0 = f32[1000, 10] parameter(0) + ROOT a0 = f32[1000, 10] exponential(p0) + } + )"); + auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root); + ASSERT_OK_AND_ASSIGN(auto original_space, + TilingSpace::Create(*fusion_adaptor, &mlir_context_)); + + mlir::MLIRContext target_context; + std::unique_ptr cloned_space = + original_space->Clone(&target_context); + ASSERT_NE(cloned_space, nullptr); + EXPECT_EQ(cloned_space->mlir_context(), &target_context); + EXPECT_EQ(cloned_space->num_dimensions(), original_space->num_dimensions()); + + for (const auto& root_tile : cloned_space->tiled_roots()) { + for (const auto& dim_tile : root_tile.dim_tiles()) { + EXPECT_EQ(dim_tile.size.GetContext(), &target_context); + EXPECT_EQ(dim_tile.offset.GetContext(), &target_context); + EXPECT_EQ(dim_tile.stride.GetContext(), &target_context); + EXPECT_EQ(dim_tile.upper_bound.GetContext(), &target_context); + } + } + + EXPECT_TRUE(cloned_space->IsSymbolic()); + EXPECT_OK(cloned_space->AssignTileSizes({64, 2})); + EXPECT_FALSE(cloned_space->IsSymbolic()); + EXPECT_EQ(cloned_space->dimensions()[0].tile_size, 64); + EXPECT_EQ(cloned_space->dimensions()[1].tile_size, 2); + EXPECT_TRUE(original_space->IsSymbolic()); +} + } // namespace } // namespace xla::gpu::experimental