From e2c7deec2444209d0201bc4c01f9207ae6873d58 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 28 Sep 2026 09:58:24 -0700 Subject: [PATCH] Squashing commits after sync tool breakage Reverts changelist 989582731 PiperOrigin-RevId: 989691756 --- .../analysis/resource_value_typed_analyzer.cc | 8 +- .../compiler/tests/bucketize_op_test.py | 9 + .../compiler/tf2xla/kernels/in_topk_op.cc | 18 +- .../core/kernels/bucketize_op_gpu.cu.cc | 10 +- tensorflow/core/kernels/topk_op.cc | 143 +- tensorflow/python/eager/BUILD | 2 +- tensorflow/python/kernel_tests/math_ops/BUILD | 4 + .../math_ops/bucketize_op_test.py | 59 +- .../kernel_tests/math_ops/in_topk_op_test.py | 18 + .../kernel_tests/math_ops/topk_op_test.py | 221 ++ tensorflow/python/ops/sort_ops_test.py | 39 + third_party/xla/build_tools/ci/BUILD | 20 + third_party/xla/build_tools/ci/bazel_diff.py | 283 +-- .../xla/build_tools/ci/bazel_diff_test.py | 188 +- third_party/xla/build_tools/ci/build.py | 64 +- .../xla/build_tools/ci/change_detector.py | 389 ++++ .../build_tools/ci/change_detector_test.py | 293 +++ .../xla/third_party/llvm/generated.patch | 705 +----- .../xla/third_party/llvm/workspace.bzl | 4 +- .../xla/third_party/shardy/temporary.patch | 1911 ++++++----------- .../xla/third_party/shardy/workspace.bzl | 4 +- .../xla/backends/autotuner/config_runner.cc | 19 +- .../backends/autotuner/config_runner_test.cc | 66 + .../backends/cpu/benchmarks/hlo/hlo_files.bzl | 8 + .../hlo/pairwise_biot_savart_4096_f64.hlo | 105 + .../pairwise_biot_savart_grad_2048_f64.hlo | 278 +++ .../hlo/pairwise_gravity_4096_f64.hlo | 53 + .../hlo/pairwise_gravity_grad_2048_f64.hlo | 147 ++ .../pairwise_min_distance_grad_2048_f64.hlo | 108 + .../hlo/pairwise_rbf_gradient_4096_f32.hlo | 48 + .../pairwise_rbf_gradient_grad_2048_f32.hlo | 115 + .../pairwise_sum_distances_grad_2048_f32.hlo | 94 + .../xla/xla/backends/cpu/codegen/BUILD | 14 + .../codegen/dot/dot_kernel_emitter_test.py | 1 - .../xla/backends/cpu/codegen/elemental/BUILD | 67 - .../elemental/concatenate_kernel_emitter.cc | 14 +- .../concatenate_kernel_emitter_test.py | 11 +- .../elemental/elemental_kernel_emitter.cc | 340 --- .../elemental/elemental_kernel_emitter.h | 68 - .../elemental_kernel_emitter_test.cc | 179 -- .../xla/backends/cpu/codegen/emitters/BUILD | 3 + .../emitters/cpu_fusion_emitter_test.cc | 250 ++- .../backends/cpu/codegen/execution_engine.cc | 6 +- .../backends/cpu/codegen/fusion_emitter.cc | 75 +- .../xla/backends/cpu/codegen/ir_compiler.cc | 12 +- .../backends/cpu/codegen/ir_compiler_test.cc | 45 +- .../cpu/codegen/object_buffer_identifier.h | 65 + .../xla/xla/backends/cpu/testlib/BUILD | 19 - .../xla/xla/backends/cpu/testlib/__init__.py | 1 - .../testlib/elemental_kernel_emitter_test.py | 391 ---- .../cpu/testlib/kernel_runner_extension.cc | 12 +- .../gpu/codegen/triton/collective_emitter.cc | 66 +- .../gpu/codegen/triton/lowering_util.cc | 22 + .../backends/gpu/codegen/triton/tests/BUILD | 1 + .../tests/collectives/all_gather_one_shot.hlo | 58 + .../triton/tests/collectives/all_reduce.hlo | 16 +- .../tests/reduce/reduce_diamond_tiled.hlo | 25 +- .../backends/gpu/codegen/triton/tools/BUILD | 35 + .../triton/tools/fusion_to_tiled_hlo.cc | 255 +++ .../triton_xla_extract_insert_to_triton.mlir | 32 + .../gpu/collectives/gpu_communicator.h | 12 + .../gpu/collectives/nccl_communicator.cc | 39 + .../gpu/collectives/nccl_communicator.h | 6 + .../libraries/native_custom_call_thunks/BUILD | 7 +- .../xla/xla/backends/gpu/runtime/BUILD | 82 +- .../runtime/all_gather_thunk_multigpu_test.cc | 27 +- .../gpu/runtime/all_reduce_build_info_test.cc | 36 +- .../runtime/all_reduce_thunk_multigpu_test.cc | 34 +- .../runtime/all_to_all_thunk_multigpu_test.cc | 24 +- ...ollective_broadcast_thunk_multigpu_test.cc | 43 +- .../gpu/runtime/collective_kernel_thunk.cc | 11 +- .../gpu/runtime/collective_kernel_thunk.proto | 6 + .../runtime/collective_kernel_thunk_test.cc | 45 +- .../backends/gpu/runtime/collective_params.h | 11 + .../collective_permute_thunk_multigpu_test.cc | 39 +- .../gpu/runtime/collective_reduce_thunk.cc | 248 +++ .../gpu/runtime/collective_reduce_thunk.h | 113 + .../collective_reduce_thunk_multigpu_test.cc | 214 ++ .../backends/gpu/runtime/collective_thunk.cc | 10 +- .../collective_thunk_multigpu_test_utils.h | 3 + .../xla/xla/backends/gpu/runtime/command.h | 31 - .../gpu/runtime/command_buffer_cmd_emitter.cc | 28 +- .../command_buffer_cmd_emitter_test.cc | 105 +- .../runtime/command_buffer_conversion_pass.cc | 3 + .../backends/gpu/runtime/command_executor.cc | 105 +- .../backends/gpu/runtime/command_executor.h | 5 +- .../gpu/runtime/command_executor_test.cc | 6 +- .../runtime/send_recv_thunk_multigpu_test.cc | 35 +- .../xla/xla/backends/gpu/runtime/thunk.cc | 9 +- .../xla/xla/backends/gpu/runtime/thunk.h | 1 + .../xla/xla/backends/gpu/runtime/thunk.proto | 10 + .../xla/backends/gpu/runtime/thunk_kind.proto | 1 + .../runtime/thunk_proto_deserialization.cc | 5 + third_party/xla/xla/backends/gpu/tests/BUILD | 3 - .../gpu/tests/collective_ops_e2e_test.cc | 115 + .../backends/gpu/tests/single_instruction.hlo | 21 +- .../xla/xla/benchmarks/core/benchmark.py | 1 + .../transforms/expand_integer_power.cc | 18 + .../emitters/transforms/lower_tensors.cc | 171 +- .../tests/expand_integer_power.mlir | 26 + .../transforms/tests/lower_tensors.mlir | 320 +-- .../codegen/xtile/codegen/emitter_helpers.h | 50 +- .../codegen/experimental_fusion_emitter.cc | 56 +- .../ir/transforms/lower_stablehlo_to_arith.cc | 3 + .../tests/lower_stablehlo_to_arith.mlir | 10 + third_party/xla/xla/debug_options_flags.cc | 10 + third_party/xla/xla/hlo/parser/hlo_parser.cc | 8 +- .../collectives/async_collective_creator.cc | 15 +- .../collectives/async_collective_creator.h | 1 + .../python/ifrt_proxy/client/executable.cc | 5 +- .../python/ifrt_proxy/client/rpc_helper.cc | 20 +- .../xla/python/ifrt_proxy/client/rpc_helper.h | 12 +- .../xla/xla/service/buffer_assignment.cc | 1 + third_party/xla/xla/service/cpu/BUILD | 27 - .../xla/xla/service/cpu/cpu_compiler.cc | 18 +- .../cpu/cpu_compiler_internals_test.cc | 10 +- .../xla/xla/service/cpu/cpu_compiler_test.cc | 16 + .../xla/service/cpu/cpu_instruction_fusion.cc | 16 +- .../cpu/cpu_instruction_fusion_test.cc | 5 - .../xla/xla/service/cpu/cpu_options.cc | 9 - third_party/xla/xla/service/cpu/cpu_options.h | 3 - .../xla/xla/service/cpu/fusion_wrapper.cc | 2 +- .../xla/xla/service/cpu/fusion_wrapper.h | 5 +- .../xla/service/cpu/fusion_wrapper_test.cc | 81 +- .../xla/xla/service/cpu/ir_emission_utils.h | 15 - .../xla/xla/service/cpu/ir_emitter2.cc | 216 -- third_party/xla/xla/service/cpu/ir_emitter2.h | 45 - .../xla/service/cpu/parallel_loop_emitter.cc | 95 - .../xla/service/cpu/parallel_loop_emitter.h | 77 - third_party/xla/xla/service/cpu/tests/BUILD | 18 - .../xla/service/cpu/tests/cpu_copy_test.cc | 14 - .../cpu/tests/cpu_external_constants_test.cc | 74 - .../xla/xla/service/cpu/thunk_emitter.cc | 145 +- .../xla/xla/service/cpu/thunk_emitter.h | 12 +- third_party/xla/xla/service/gpu/BUILD | 1 + .../xla/xla/service/gpu/gpu_hlo_schedule.cc | 8 + .../xla/xla/service/gpu/thunk_emitter.cc | 37 +- .../stream_executor/gpu/all_reduce_kernel.h | 2 +- .../stream_executor/sycl/sycl_gpu_runtime.cc | 5 +- third_party/xla/xla/tests/BUILD | 1 + .../aot_compatibility_experimental/gpu/BUILD | 1 - .../xla/tests/aot_interception_pjrt_client.cc | 10 + third_party/xla/xla/util/split_proto/BUILD | 1 + third_party/xla/xla/xla.proto | 9 +- 144 files changed, 6005 insertions(+), 4709 deletions(-) create mode 100644 third_party/xla/build_tools/ci/change_detector.py create mode 100644 third_party/xla/build_tools/ci/change_detector_test.py create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_4096_f64.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_grad_2048_f64.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_4096_f64.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_grad_2048_f64.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_min_distance_grad_2048_f64.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_4096_f32.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_grad_2048_f32.hlo create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_sum_distances_grad_2048_f32.hlo delete mode 100644 third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.cc delete mode 100644 third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h delete mode 100644 third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter_test.cc create mode 100644 third_party/xla/xla/backends/cpu/codegen/object_buffer_identifier.h delete mode 100644 third_party/xla/xla/backends/cpu/testlib/elemental_kernel_emitter_test.py create mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo create mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tools/fusion_to_tiled_hlo.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc delete mode 100644 third_party/xla/xla/service/cpu/parallel_loop_emitter.cc delete mode 100644 third_party/xla/xla/service/cpu/parallel_loop_emitter.h delete mode 100644 third_party/xla/xla/service/cpu/tests/cpu_external_constants_test.cc diff --git a/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc b/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc index 29b93f10e83950..364be1bd659729 100644 --- a/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc +++ b/tensorflow/compiler/mlir/tensorflow/analysis/resource_value_typed_analyzer.cc @@ -148,10 +148,10 @@ LogicalResult ResourceAnalyzer::AnalyzeRegion(Region& region) { return; } if (auto if_op = dyn_cast(op)) { - PropagatePotentiallyWrittenUpFromCallee(if_op.getThenBranch(), - if_op.getODSOperands(1)); - PropagatePotentiallyWrittenUpFromCallee(if_op.getElseBranch(), - if_op.getODSOperands(1)); + PropagatePotentiallyWrittenUpFromCallee( + if_op.getThenBranch(), if_op->getOperands().drop_front()); + PropagatePotentiallyWrittenUpFromCallee( + if_op.getElseBranch(), if_op->getOperands().drop_front()); return; } if (auto while_op = dyn_cast(op)) { diff --git a/tensorflow/compiler/tests/bucketize_op_test.py b/tensorflow/compiler/tests/bucketize_op_test.py index d6891febdd000f..c8fce3435ac7e0 100644 --- a/tensorflow/compiler/tests/bucketize_op_test.py +++ b/tensorflow/compiler/tests/bucketize_op_test.py @@ -34,6 +34,15 @@ def testInt(self): self.assertAllEqual(expected_out, sess.run(op, {p: [-5, 0, 2, 3, 5, 8, 10, 11, 12]})) + def testIntWithFloatBoundaries(self): + for dtype in [dtypes.int32, dtypes.int64]: + with self.session() as sess: + p = array_ops.placeholder(dtype) + with self.test_scope(): + op = math_ops._bucketize(p, boundaries=[0.1, 1.1]) + expected_out = [0, 1, 2, 2] + self.assertAllEqual(expected_out, sess.run(op, {p: [0, 1, 2, 3]})) + def testFloat(self): with self.session() as sess: p = array_ops.placeholder(dtypes.float32) diff --git a/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc b/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc index e22a902853a528..57a0f2562a53f8 100644 --- a/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/in_topk_op.cc @@ -28,8 +28,6 @@ limitations under the License. #include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" -#include "tensorflow/core/platform/errors.h" -#include "tensorflow/core/platform/types.h" namespace tensorflow { namespace { @@ -79,6 +77,17 @@ class InTopKOp : public XlaOpKernel { xla::XlaOp iota_r2 = xla::Broadcast(iota_r1, {batch_size}); xla::XlaOp eq_r2 = xla::Eq(targets_r1, iota_r2, {0}); + // The masked reduction below treats an out-of-range target as zero, which + // can otherwise produce a false hit. Check bounds per target, rather than + // reducing the class-match matrix again. + xla::XlaOp targets_s64 = xla::ConvertElementType(targets_r1, xla::S64); + // This folds to a constant for static dimensions, while using the actual + // class count instead of the static bound for a dynamic dimension. + xla::XlaOp num_classes_s64 = xla::ConvertElementType( + xla::GetDimensionSize(predictions_r2, 1), xla::S64); + xla::XlaOp valid_target_r1 = + xla::And(xla::Ge(targets_s64, xla::ConstantR0(xla_builder, 0)), + xla::Lt(targets_s64, num_classes_s64)); xla::XlaOp zero_r0_f32 = xla::Zero(xla_builder, xla::F32); xla::XlaOp zero_r2_f32 = xla::ZerosLike(predictions_r2); xla::XlaOp select_r2 = xla::Select(eq_r2, predictions_r2, zero_r2_f32); @@ -99,9 +108,10 @@ class InTopKOp : public XlaOpKernel { one_hot_r2, zero_r0, xla::CreateScalarAddComputation(xla::S32, xla_builder), {1}); - xla::XlaOp result = + xla::XlaOp result = xla::And( xla::And(xla::Lt(num_gt_r1, xla::ConstantR0(xla_builder, k)), - xla::IsFinite(targets_values_r1)); + xla::IsFinite(targets_values_r1)), + valid_target_r1); context->SetOutput(0, result); } diff --git a/tensorflow/core/kernels/bucketize_op_gpu.cu.cc b/tensorflow/core/kernels/bucketize_op_gpu.cu.cc index 2b0429d4a67cb3..125f9dd2453ade 100644 --- a/tensorflow/core/kernels/bucketize_op_gpu.cu.cc +++ b/tensorflow/core/kernels/bucketize_op_gpu.cu.cc @@ -17,6 +17,8 @@ limitations under the License. #define EIGEN_USE_GPU +#include + #include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" @@ -55,15 +57,21 @@ __global__ void BucketizeCustomKernel( boundaries = shared_mem_boundaries; } + // For double input, compare as double. For float and integer inputs, compare + // as float to match CPU and XLA behavior without truncating fractional + // boundaries. + using CompType = std::conditional_t, double, float>; + GPU_1D_KERNEL_LOOP(i, size_in) { T value = in[i]; + CompType value_cast = static_cast(value); int32_t bucket = 0; int32_t count = size_boundaries; while (count > 0) { int32_t l = bucket; int32_t step = count / 2; l += step; - if (!(value < static_cast(boundaries[l]))) { + if (!(value_cast < static_cast(boundaries[l]))) { bucket = ++l; count -= step + 1; } else { diff --git a/tensorflow/core/kernels/topk_op.cc b/tensorflow/core/kernels/topk_op.cc index a1ad47d265bc16..26943126782793 100644 --- a/tensorflow/core/kernels/topk_op.cc +++ b/tensorflow/core/kernels/topk_op.cc @@ -20,6 +20,10 @@ limitations under the License. #include "tensorflow/core/kernels/topk_op.h" #include +#include +#include +#include +#include #include #include @@ -30,6 +34,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/top_n.h" +#include "tensorflow/core/platform/macros.h" // NOLINT(misc-include-cleaner) #include "tensorflow/core/util/work_sharder.h" namespace tensorflow { @@ -132,6 +137,24 @@ class TopK : public OpKernel { namespace functor { +namespace { + +// Returns +1 for +NaN, -1 for -NaN, and 0 for any other value. Ordering by +// value and then by this rank gives +NaN > +Inf > ... > -Inf > -NaN, which +// matches the GPU radix-sort key order and the XLA total order for NaNs. +// +0.0 and -0.0 compare equal and are ordered by index. +template +EIGEN_ALWAYS_INLINE int NanRank(const T v) { // NOLINT(misc-include-cleaner) + // NOLINTNEXTLINE(misc-include-cleaner) + if (TF_PREDICT_TRUE(!Eigen::numext::isnan(v))) return 0; + // Eigen::numext::signbit returns a Scalar bit mask (all ones, itself a NaN + // for floating types), not a bool, so test the sign with std::signbit. The + // conversion to double preserves the sign bit of a NaN for every float type. + return std::signbit(static_cast(v)) ? -1 : 1; +} + +} // namespace + template struct TopKFunctor { static EIGEN_ALWAYS_INLINE absl::Status Compute( @@ -139,46 +162,51 @@ struct TopKFunctor { const typename TTypes::ConstTensor& input, const int64_t num_rows, const int64_t num_cols, typename TTypes::Tensor values, typename TTypes::Tensor indices) { - const CPUDevice& d = context->eigen_device(); - - // Special case for k == 1. - if (k == 1) { - typename Eigen::IndexList> reduce_on_cols; - typename Eigen::IndexList> rows_by_one; - rows_by_one.set(0, num_rows); - - values.device(d) = - input.maximum(/*dims=*/reduce_on_cols).eval().reshape(rows_by_one); - // Get the indices of the maximum values. - for (int r = 0; r < num_rows; ++r) { - indices(r, 0) = Tidx(0); - for (int c = 0; c < num_cols; ++c) { - if (values(r, 0) == input(r, c)) { - indices(r, 0) = static_cast(c); - break; - } - } - values(r, 0) = input(r, indices(r, 0)); - } - - return absl::OkStatus(); - } - auto SortIndices = [&](int64_t start_batch, int64_t limit_batch) { - for (int32_t b = start_batch; b < limit_batch; ++b) { + for (int64_t b = start_batch; b < limit_batch; ++b) { const T* input_data = &input(b, 0); - const auto stable_comp = [input_data](const int32_t a, - const int32_t b) { - if (input_data[b] < input_data[a]) { + // Both comparators order values descending with the relational + // comparisons first (the common path). Equal values (never NaN, and + // including +0/-0) tie; only a NaN operand falls back to the NaN rank, + // so NaNs are ordered +NaN first and -NaN last. + const auto stable_comp = [input_data](const Tidx a, const Tidx b) { + const T val_a = input_data[a]; + const T val_b = input_data[b]; + if (val_b < val_a) { return true; - } else if (input_data[b] > input_data[a]) { + } + if (val_a < val_b) { return false; - } else { + } + if (val_a == val_b) { return a < b; } + if constexpr (!Eigen::NumTraits::IsInteger) { + const int rank_a = NanRank(val_a); + const int rank_b = NanRank(val_b); + if (rank_a != rank_b) { + return rank_b < rank_a; + } + } + return a < b; }; - const auto comp = [input_data](const int32_t a, const int32_t b) { - return input_data[b] < input_data[a]; + const auto comp = [input_data](const Tidx a, const Tidx b) { + const T val_a = input_data[a]; + const T val_b = input_data[b]; + if (val_b < val_a) { + return true; + } + if (val_a < val_b) { + return false; + } + if (val_a == val_b) { + return false; + } + if constexpr (!Eigen::NumTraits::IsInteger) { + return NanRank(val_b) < NanRank(val_a); + } else { + return false; + } }; // TODO(ebrevdo): For large k < num_cols, instead of using // TopN, it may be faster to create a temporary vector of @@ -194,19 +222,50 @@ struct TopKFunctor { // indices that started out sorted. First, do a std::sort, which // is notably faster than std::stable_sort. std::sort(begin, end, comp); - // Then, for runs of adjacent elements that were equal, sort the - // indices in those runs in increasing order. + // Then, for runs of adjacent elements that are equivalent under + // `comp` (equal values, or NaNs of the same sign), sort the indices + // in those runs in increasing order. for (auto* run_begin = begin; run_begin != end;) { auto* run_end = run_begin + 1; - if (run_end == end) break; - if (input_data[*run_begin] == input_data[*run_end]) { - while (++run_end != end) { - if (input_data[*run_begin] != input_data[*run_end]) break; + while (run_end != end && !comp(*run_begin, *run_end)) ++run_end; + if (run_end - run_begin > 1) std::sort(run_begin, run_end); + run_begin = run_end; + } + } else if (k == 1) { + // A single scan instead of an Eigen `maximum` reduction: the NaN + // propagation of the reduction is unspecified, and locating its + // result by equality can never find a NaN. Returns the first +NaN + // if any, else the first maximum non-NaN value, else the first -NaN. + Tidx best_c = 0; + T best_val = input_data[0]; + Tidx first_c = 1; + if constexpr (!Eigen::NumTraits::IsInteger) { + // A leading +NaN is the maximum: skip the scan. + if (NanRank(best_val) > 0) first_c = static_cast(num_cols); + } + for (Tidx c = first_c; c < num_cols; ++c) { + const T val = input_data[c]; + // '<=' is false only if val > best_val or either operand is NaN. + if (val <= best_val) continue; + if (best_val < val) { + best_val = val; + best_c = c; + continue; + } + if constexpr (!Eigen::NumTraits::IsInteger) { + const int rank = NanRank(val); + if (rank > 0) { // +NaN: the global maximum, stop. + best_c = c; + break; + } + if (rank == 0) { // best_val is -NaN; any non-NaN beats it. + best_val = val; + best_c = c; } - std::sort(run_begin, run_end); + // rank < 0: a -NaN never beats the current best. } - run_begin = run_end; } + indices(b, 0) = best_c; } else { // Use the TopN heap object to sort. gtl::TopN filter(k, stable_comp); @@ -215,7 +274,7 @@ struct TopKFunctor { filter.push(c); } - int32_t i = 0; + Tidx i = 0; if (sorted) { std::unique_ptr> top_k(filter.Extract()); for (auto top_k_it = top_k->begin(); top_k_it != top_k->end(); diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD index b35b844ba2d344..64b065c926fa6b 100644 --- a/tensorflow/python/eager/BUILD +++ b/tensorflow/python/eager/BUILD @@ -1151,7 +1151,7 @@ cuda_py_strict_test( srcs = ["remote_test.py"], exec_properties = select({ # copybara:uncomment_begin(google-only) - # "//tools/cpp:msan_build": {"mem": "30g"}, + # "//tools/cpp:msan_build": {"mem": "28g"}, # copybara:uncomment_end "//conditions:default": None, }), diff --git a/tensorflow/python/kernel_tests/math_ops/BUILD b/tensorflow/python/kernel_tests/math_ops/BUILD index 0edf347a193793..97a290b2f20646 100644 --- a/tensorflow/python/kernel_tests/math_ops/BUILD +++ b/tensorflow/python/kernel_tests/math_ops/BUILD @@ -360,7 +360,11 @@ cuda_py_strict_test( size = "small", srcs = ["in_topk_op_test.py"], deps = [ + "//tensorflow/compiler/jit:xla_cpu_jit", + "//tensorflow/compiler/jit:xla_gpu_jit", + "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:constant_op", + "//tensorflow/python/framework:dtypes", "//tensorflow/python/ops:nn_ops", "//tensorflow/python/platform:client_testlib", "//third_party/py/numpy", diff --git a/tensorflow/python/kernel_tests/math_ops/bucketize_op_test.py b/tensorflow/python/kernel_tests/math_ops/bucketize_op_test.py index 6f1003c8d1b11b..b6a709f40b12e7 100644 --- a/tensorflow/python/kernel_tests/math_ops/bucketize_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/bucketize_op_test.py @@ -32,14 +32,53 @@ def testInt(self): constant_op.constant([-5, 0, 2, 3, 5, 8, 10, 11, 12]), boundaries=[0, 3, 8, 11]) expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4] - with self.session(): + with self.session(use_gpu=True): + self.assertAllEqual(expected_out, self.evaluate(op)) + + def testInt64(self): + op = math_ops._bucketize( + constant_op.constant( + [-5, 0, 2, 3, 5, 8, 10, 11, 12], dtype=dtypes.int64 + ), + boundaries=[0, 3, 8, 11], + ) + expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4] + with self.session(use_gpu=True): self.assertAllEqual(expected_out, self.evaluate(op)) + def testIntWithFloatBoundaries(self): + for dtype in [dtypes.int32, dtypes.int64]: + # Positive fractional boundaries (reproducer from issue #70343) + op = math_ops._bucketize( + constant_op.constant([0, 1, 2, 3], dtype=dtype), boundaries=[0.1, 1.1] + ) + expected_out = [0, 1, 2, 2] + with self.session(use_gpu=True): + self.assertAllEqual(expected_out, self.evaluate(op)) + + # Negative and positive fractional boundaries + op = math_ops._bucketize( + constant_op.constant([-2, -1, 0, 1, 2], dtype=dtype), + boundaries=[-1.5, -0.5, 0.5, 1.5], + ) + expected_out = [0, 1, 2, 3, 4] + with self.session(use_gpu=True): + self.assertAllEqual(expected_out, self.evaluate(op)) + + # Fractional boundaries close to integer boundaries + op = math_ops._bucketize( + constant_op.constant([0, 1, 2], dtype=dtype), + boundaries=[0.999, 1.001], + ) + expected_out = [0, 1, 2] + with self.session(use_gpu=True): + self.assertAllEqual(expected_out, self.evaluate(op)) + def testEmptyFloat(self): op = math_ops._bucketize( array_ops.zeros([0, 3], dtype=dtypes.float32), boundaries=[]) expected_out = np.zeros([0, 3], dtype=np.float32) - with self.session(): + with self.session(use_gpu=True): self.assertAllEqual(expected_out, self.evaluate(op)) def testFloat(self): @@ -47,7 +86,19 @@ def testFloat(self): constant_op.constant([-5., 0., 2., 3., 5., 8., 10., 11., 12.]), boundaries=[0., 3., 8., 11.]) expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4] - with self.session(): + with self.session(use_gpu=True): + self.assertAllEqual(expected_out, self.evaluate(op)) + + def testDouble(self): + op = math_ops._bucketize( + constant_op.constant( + [-5.0, 0.0, 2.0, 3.0, 5.0, 8.0, 10.0, 11.0, 12.0], + dtype=dtypes.float64, + ), + boundaries=[0.0, 3.0, 8.0, 11.0], + ) + expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4] + with self.session(use_gpu=True): self.assertAllEqual(expected_out, self.evaluate(op)) def test2DInput(self): @@ -55,7 +106,7 @@ def test2DInput(self): constant_op.constant([[-5, 0, 2, 3, 5], [8, 10, 11, 12, 0]]), boundaries=[0, 3, 8, 11]) expected_out = [[0, 1, 1, 2, 2], [3, 3, 4, 4, 1]] - with self.session(): + with self.session(use_gpu=True): self.assertAllEqual(expected_out, self.evaluate(op)) @test_util.run_deprecated_v1 diff --git a/tensorflow/python/kernel_tests/math_ops/in_topk_op_test.py b/tensorflow/python/kernel_tests/math_ops/in_topk_op_test.py index fc8e336cfe26e9..27d0050119fd8b 100644 --- a/tensorflow/python/kernel_tests/math_ops/in_topk_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/in_topk_op_test.py @@ -16,7 +16,9 @@ import numpy as np +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes from tensorflow.python.ops import nn_ops from tensorflow.python.platform import test @@ -62,6 +64,22 @@ def testBadTarget(self): target = [2, 4] # must return False for invalid target self._validateInTopK(predictions, target, 2, [True, False]) + def testXlaOutOfRangeTarget(self): + predictions = constant_op.constant( + [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]] + ) + expected = [False, True, True, False] + for dtype in (dtypes.int32, dtypes.int64): + targets = constant_op.constant([-1, 0, 1, 2], dtype=dtype) + + @def_function.function(jit_compile=True) + def in_top_k(predictions, targets): + return nn_ops.in_top_k(predictions, targets, 1) + + self.assertAllEqual( + self.evaluate(in_top_k(predictions, targets)), expected + ) + def testEmpty(self): predictions = np.empty([0, 5]) target = np.empty([0], np.int32) diff --git a/tensorflow/python/kernel_tests/math_ops/topk_op_test.py b/tensorflow/python/kernel_tests/math_ops/topk_op_test.py index 61c35f17ac9b5c..c9a5cae24236d4 100644 --- a/tensorflow/python/kernel_tests/math_ops/topk_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/topk_op_test.py @@ -139,6 +139,227 @@ def testTop1AllNan(self): inputs = [[np.nan, np.nan], [np.nan, np.nan]] self._validateTopK(inputs, 1, [[np.nan], [np.nan]], [[0], [0]]) + def testTop1WithNan(self) -> None: + neg_nan = np.copysign(np.nan, -1.0) + inputs = [ + [3.0, np.nan, 2.0], + [1.0, 2.0, 3.0, np.nan], + [np.nan, 5.0, np.nan], + [1.0, 5.0, 2.0], + # A leading +NaN is the maximum and short-circuits the scan. + [np.nan, 1000.0, np.nan, 2.0], + # A leading -NaN does not: the scan continues to the later +NaN. + [neg_nan, 1.0, np.nan, 2.0], + ] + expected_indices = [1, 3, 0, 1, 0, 2] + with self.cached_session(): + for row, expected_idx in zip(inputs, expected_indices): + values, indices = self.evaluate(nn_ops.top_k(row, 1)) + self.assertAllEqual(indices, [expected_idx]) + if np.isnan(row[expected_idx]): + self.assertTrue(np.isnan(values[0])) + else: + self.assertEqual(values[0], row[expected_idx]) + + # The two leading-NaN rows again, batched with a finite row. + batched = np.array( + inputs[4:] + [[4.0, 3.0, 1000.0, 2.0]], dtype=np.float32 + ) + values, indices = self.evaluate(nn_ops.top_k(batched, 1)) + self.assertAllEqual(indices, [[0], [2], [2]]) + self.assertTrue(np.all(np.isnan(values[:2]))) + self.assertFalse(np.any(np.signbit(values[:2]))) + self.assertEqual(values[2, 0], 1000.0) + + def testTopKNanEdgeCases(self) -> None: + # +NaN sorts first and -NaN last, as on GPU (radix sort on the key bits) + # and in the XLA total order. + neg_nan = np.copysign(np.nan, -1.0) + pos_nan = np.nan + inputs = [-np.inf, pos_nan, 0.0, np.inf, -0.0, neg_nan] + with self.cached_session(): + values, indices = self.evaluate(nn_ops.top_k(inputs, 6)) + self.assertAllEqual(indices, [1, 3, 2, 4, 0, 5]) + self.assertTrue(np.isnan(values[0])) + self.assertAllEqual(values[1:5], [np.inf, 0.0, -0.0, -np.inf]) + self.assertTrue(np.isnan(values[5])) + + all_nans = [pos_nan, neg_nan, pos_nan, neg_nan] + for k, expected in ((1, [0]), (2, [0, 2]), (4, [0, 2, 1, 3])): + vals, idxs = self.evaluate(nn_ops.top_k(all_nans, k)) + self.assertTrue(np.all(np.isnan(vals))) + self.assertAllEqual(idxs, expected) + + # k == 1 prefers any non-NaN value over -NaN. + _, idxs = self.evaluate(nn_ops.top_k([neg_nan, -np.inf, neg_nan], 1)) + self.assertAllEqual(idxs, [1]) + + def testTopKNegatedNan(self) -> None: + # Ascending sort is implemented as top_k(-x): negated NaNs must sort last. + inputs = -np.array( + [np.nan, 3.0, 1.0, np.nan, 2.0, np.nan, 0.5], dtype=np.float32 + ) + self.assertTrue(np.all(np.signbit(inputs[[0, 3, 5]]))) + with self.cached_session(): + values, indices = self.evaluate(nn_ops.top_k(inputs, 7)) + self.assertAllEqual(indices, [6, 2, 4, 1, 0, 3, 5]) + self.assertAllEqual(values[:4], [-0.5, -1.0, -2.0, -3.0]) + self.assertTrue(np.all(np.isnan(values[4:]))) + + def testTopKIntegerTypes(self) -> None: + inputs = np.array([[5, 2, 8, 8, 1], [9, 3, 7, 0, 4]], dtype=np.int32) + with self.cached_session(): + values, indices = self.evaluate(nn_ops.top_k(inputs, 1)) + self.assertAllEqual(values, [[8], [9]]) + self.assertAllEqual(indices, [[2], [0]]) + + values, indices = self.evaluate(nn_ops.top_k(inputs, 3)) + self.assertAllEqual(values, [[8, 8, 5], [9, 7, 4]]) + self.assertAllEqual(indices, [[2, 3, 0], [0, 2, 4]]) + + values, indices = self.evaluate(nn_ops.top_k(inputs, 5)) + self.assertAllEqual(values, [[8, 8, 5, 2, 1], [9, 7, 4, 3, 0]]) + self.assertAllEqual(indices, [[2, 3, 0, 1, 4], [0, 2, 4, 1, 3]]) + + # k == 1 keeps the lowest index on ties for other integer widths too, + # including a tie after the running maximum has been updated. + for dtype in ( + np.int8, + np.uint8, + np.int16, + np.uint16, + np.uint32, + np.int64, + np.uint64, + ): + with self.subTest(dtype=dtype): + info = np.iinfo(dtype) + row = np.array([[info.min, info.max, 0, info.max]], dtype=dtype) + values, indices = self.evaluate(nn_ops.top_k(row, 1)) + self.assertAllEqual(indices, [[1]]) + self.assertAllEqual(values, [[info.max]]) + + def testTopKNanDtypes(self) -> None: + dtypes_to_test = ( + np.float64, + np.float32, + np.float16, + dtypes.bfloat16.as_numpy_dtype, + ) + index_types = (dtypes.int16, dtypes.int32, dtypes.int64) + for dtype, index_type in itertools.product(dtypes_to_test, index_types): + inputs = np.array( + [np.nan, 3.0, 1.0, np.nan, 2.0, np.nan, 0.5], dtype=dtype + ) + with self.subTest(dtype=dtype, index_type=index_type): + with self.cached_session(): + for k, expected in ( + (1, [0]), + (3, [0, 3, 5]), + (7, [0, 3, 5, 1, 4, 2, 6]), + ): + values, indices = self.evaluate( + nn_ops.top_k(inputs, k, index_type=index_type) + ) + self.assertEqual(indices.dtype, index_type.as_numpy_dtype) + self.assertAllEqual(indices, expected) + self.assertTrue(np.all(np.isnan(values[:3]))) + self.assertAllClose( + values[3:], [3.0, 2.0, 1.0, 0.5][: max(k - 3, 0)] + ) + + # +NaN first and -NaN last on the k == 1, heap, and full-sort paths; + # a leading -NaN must not end the k == 1 scan early. + neg_nan_row = np.array( + [np.copysign(np.nan, -1.0), 1.0, np.nan, 2.0], dtype=dtype + ) + for k, expected in ((1, [2]), (2, [2, 3]), (4, [2, 3, 1, 0])): + values, indices = self.evaluate( + nn_ops.top_k(neg_nan_row, k, index_type=index_type) + ) + self.assertAllEqual(indices, expected) + values = values.astype(np.float32) + self.assertTrue(np.isnan(values[0])) + self.assertFalse(np.signbit(values[0])) + if k == neg_nan_row.size: # Full sort: -NaN, with its sign, last. + self.assertTrue(np.isnan(values[-1])) + self.assertTrue(np.signbit(values[-1])) + + def testTopKBatchedWithNan(self) -> None: + inputs = np.array( + [ + [3.0, np.nan, 2.0, 1.0], + [np.nan, 5.0, np.nan, 4.0], + [1.0, 2.0, 3.0, np.nan], + [4.0, 3.0, 2.0, 1.0], + ], + dtype=np.float32, + ) + with self.cached_session(): + values, indices = self.evaluate(nn_ops.top_k(inputs, 1)) + self.assertTrue(np.isnan(values[0, 0])) + self.assertEqual(indices[0, 0], 1) + self.assertTrue(np.isnan(values[1, 0])) + self.assertEqual(indices[1, 0], 0) + self.assertTrue(np.isnan(values[2, 0])) + self.assertEqual(indices[2, 0], 3) + self.assertEqual(values[3, 0], 4.0) + self.assertEqual(indices[3, 0], 0) + + values, indices = self.evaluate(nn_ops.top_k(inputs, 2)) + self.assertTrue(np.isnan(values[0, 0])) + self.assertEqual(indices[0, 0], 1) + self.assertEqual(values[0, 1], 3.0) + self.assertEqual(indices[0, 1], 0) + + self.assertTrue(np.isnan(values[1, 0])) + self.assertTrue(np.isnan(values[1, 1])) + self.assertEqual(indices[1, 0], 0) + self.assertEqual(indices[1, 1], 2) + + def testTopKLargeBatchWithNanMatchesReference(self) -> None: + # Large enough to be split across several CPU shards. Pinned to CPU: the + # GPU heap path used for these shapes is not covered by this change. + rng = np.random.default_rng(1234) + num_rows, num_cols = 256, 2048 + # Quarter steps create many ties (including -0.0 vs 0.0). + inputs = ( + np.round(rng.standard_normal((num_rows, num_cols)) * 4) / 4 + ).astype(np.float32) + nan_mask = rng.random((num_rows, num_cols)) < 0.05 + signs = np.where(rng.random((num_rows, num_cols)) < 0.5, -1.0, 1.0) + inputs[nan_mask] = np.copysign(np.nan, signs[nan_mask]) + + # Reference order: +NaN, then values descending, then -NaN; ties by index. + is_nan = np.isnan(inputs) + group = np.where(is_nan, np.where(np.signbit(inputs), 2, 0), 1) + neg_values = np.where(is_nan, 0.0, -inputs) + col = np.broadcast_to(np.arange(num_cols), inputs.shape) + expected = np.stack( + [np.lexsort((col[r], neg_values[r], group[r])) for r in range(num_rows)] + ) + + with self.cached_session(), ops.device("/cpu:0"): + for k in (1, 37, num_cols): + _, indices = self.evaluate(nn_ops.top_k(inputs, k)) + self.assertAllEqual(indices, expected[:, :k]) + + def testTopKNanUnsorted(self) -> None: + inputs = [np.nan, 3.0, 1.0, np.nan, 2.0, np.nan, 0.5] + with self.cached_session(): + values, indices = self.evaluate(nn_ops.top_k(inputs, 3, sorted=False)) + self.assertTrue(np.all(np.isnan(values))) + self.assertEqual(set(indices), {0, 3, 5}) + + # k == 1 and k == num_cols take dedicated code paths. + values, indices = self.evaluate(nn_ops.top_k(inputs, 1, sorted=False)) + self.assertTrue(np.isnan(values[0])) + self.assertAllEqual(indices, [0]) + + values, indices = self.evaluate(nn_ops.top_k(inputs, 7, sorted=False)) + self.assertEqual(sorted(indices), list(range(7))) + self.assertEqual(np.sum(np.isnan(values)), 3) + def _testLargeSort(self, dtype): b = 10 n = 5000 diff --git a/tensorflow/python/ops/sort_ops_test.py b/tensorflow/python/ops/sort_ops_test.py index 77829e29f245dd..7116a33e803082 100644 --- a/tensorflow/python/ops/sort_ops_test.py +++ b/tensorflow/python/ops/sort_ops_test.py @@ -181,6 +181,45 @@ def testArgsortTensorShape(self): placeholder.shape.as_list(), sort_ops.argsort(placeholder, axis=axis).shape.as_list()) + def testSortWithNaN(self) -> None: + arr = constant_op.constant( + [np.nan, 3.0, 1.0, np.nan, 2.0, np.nan, 0.5], dtype=dtypes.float32 + ) + with self.cached_session(): + sorted_asc = self.evaluate(sort_ops.sort(arr, direction='ASCENDING')) + sorted_desc = self.evaluate(sort_ops.sort(arr, direction='DESCENDING')) + argsort_asc = self.evaluate(sort_ops.argsort(arr, direction='ASCENDING')) + argsort_desc = self.evaluate( + sort_ops.argsort(arr, direction='DESCENDING') + ) + + # NaNs go last in ascending order and first in descending order. + self.assertAllClose(sorted_asc[:4], [0.5, 1.0, 2.0, 3.0]) + self.assertTrue(np.all(np.isnan(sorted_asc[4:]))) + self.assertAllEqual(argsort_asc, [6, 2, 4, 1, 0, 3, 5]) + + self.assertTrue(np.all(np.isnan(sorted_desc[:3]))) + self.assertAllClose(sorted_desc[3:], [3.0, 2.0, 1.0, 0.5]) + self.assertAllEqual(argsort_desc, [0, 3, 5, 1, 4, 2, 6]) + + def testSortWithNegativeNaNAndInf(self) -> None: + # +NaN sorts above +Inf and -NaN below -Inf, as on GPU. + neg_nan = np.copysign(np.nan, -1.0) + pos_nan = np.nan + arr = constant_op.constant( + [-np.inf, pos_nan, 0.0, np.inf, -0.0, neg_nan], dtype=dtypes.float32 + ) + with self.cached_session(): + sorted_desc = self.evaluate(sort_ops.sort(arr, direction='DESCENDING')) + argsort_desc = self.evaluate( + sort_ops.argsort(arr, direction='DESCENDING') + ) + + self.assertTrue(np.isnan(sorted_desc[0])) + self.assertAllEqual(sorted_desc[1:5], [np.inf, 0.0, -0.0, -np.inf]) + self.assertTrue(np.isnan(sorted_desc[5])) + self.assertAllEqual(argsort_desc, [1, 3, 2, 4, 0, 5]) + if __name__ == '__main__': test.main() diff --git a/third_party/xla/build_tools/ci/BUILD b/third_party/xla/build_tools/ci/BUILD index 2cb0a7e2b7e89a..03dfa81decf83d 100644 --- a/third_party/xla/build_tools/ci/BUILD +++ b/third_party/xla/build_tools/ci/BUILD @@ -33,9 +33,27 @@ package( licenses = ["notice"], ) +pytype_strict_library( + name = "change_detector", + srcs = ["change_detector.py"], +) + +py_test( + name = "change_detector_test", + srcs = ["change_detector_test.py"], + deps = [ + ":change_detector", + "@absl_py//absl/testing:absltest", + "@absl_py//absl/testing:parameterized", + ], +) + pytype_strict_library( name = "bazel_diff", srcs = ["bazel_diff.py"], + deps = [ + ":change_detector", + ], ) py_test( @@ -43,6 +61,7 @@ py_test( srcs = ["bazel_diff_test.py"], deps = [ ":bazel_diff", + ":change_detector", "@absl_py//absl/testing:absltest", "@absl_py//absl/testing:parameterized", ], @@ -53,6 +72,7 @@ pytype_strict_binary( srcs = ["build.py"], deps = [ ":bazel_diff", + ":change_detector", ], ) diff --git a/third_party/xla/build_tools/ci/bazel_diff.py b/third_party/xla/build_tools/ci/bazel_diff.py index ad82dd0567bebe..d7cb221508db39 100644 --- a/third_party/xla/build_tools/ci/bazel_diff.py +++ b/third_party/xla/build_tools/ci/bazel_diff.py @@ -25,42 +25,20 @@ import hashlib import logging import os -import re import subprocess import sys import tempfile import time from typing import Any, List, Optional, Protocol, Tuple +from build_tools.ci import change_detector + BAZEL_DIFF_VERSION = "v46.1.0" BAZEL_DIFF_JAR_URL = "https://github.com/Tinder/bazel-diff/releases/download/v46.1.0/bazel-diff_deploy.jar" BAZEL_DIFF_JAR_SHA256 = ( "1026c66304d262e26065fc0ccd0281766c4c8b63abd18ea1fc4366fe05453c1e" ) -# Global configuration files whose modification invalidates incremental diffs -# and necessitates a full test suite run. -GLOBAL_BAZEL_CONFIG_PATTERNS: Tuple[str, ...] = ( - r"(^|/)MODULE\.bazel$", - r"(^|/)REPO\.bazel$", - r"(^|/)WORKSPACE(\.bzlmod)?$", - r"\.bazelrc$", - r"\.bazelversion$", -) - -# b/519689071: keywords in BUILD/.bzl changes that signal unconfigured hashing -# under-selection (label_flag, config_setting, alias). -FULL_RUN_BUILD_KEYWORDS: Tuple[str, ...] = ( - "label_flag(", - "config_setting(", - "alias(", -) - -# Regex matching files that only contain documentation or repository metadata. -DOCS_OR_METADATA_PATTERN = re.compile( - r"(\.md$|^docs/|(^|/)OWNERS$|^LICENSE|^\.clang|^\.gitignore|^\.vscode/|\.png$|\.jpg$|\.jpeg$|\.svg$|\.webp$|\.gif$)" -) - # Target prefixes to keep in XLA's CI. KEEP_TARGET_PREFIXES: Tuple[str, ...] = ( "//xla/", @@ -118,144 +96,6 @@ class BazelDiffDecision: elapsed_seconds: float = 0.0 -def matches_any_pattern(path: str, patterns: Sequence[str]) -> bool: - """Returns True if path matches any regex pattern in patterns.""" - return any(re.search(pattern, path) for pattern in patterns) - - -def is_global_config_changed(changed_files: Sequence[str]) -> bool: - """Returns True if any changed file matches global bazel config patterns.""" - return any( - matches_any_pattern(file_path, GLOBAL_BAZEL_CONFIG_PATTERNS) - for file_path in changed_files - ) - - -def is_docs_or_metadata_only(changed_files: Sequence[str]) -> bool: - """Returns True if all changed files are docs or non-build metadata.""" - if not changed_files: - return False - return all( - DOCS_OR_METADATA_PATTERN.search(file_path) for file_path in changed_files - ) - - -def _run_git_best_effort( - args: Sequence[str], cwd: str = "." -) -> Optional[subprocess.CompletedProcess[str]]: - """Runs a git command best-effort, returning None on OS/subprocess error.""" - try: - return subprocess.run( - ["git", *args], - cwd=cwd, - capture_output=True, - text=True, - check=False, - ) - except (OSError, subprocess.SubprocessError): - return None - - -def get_merge_base( - base_sha: str, head_sha: str = "HEAD", cwd: str = "." -) -> Optional[str]: - """Computes git merge-base between base_sha and head_sha. - - Args: - base_sha: Base git commit SHA. - head_sha: Head git commit SHA or ref. - cwd: Directory where the git command should be executed. - - Returns: - The merge base commit SHA as a string, or None if not found. - """ - result = _run_git_best_effort(["merge-base", base_sha, head_sha], cwd=cwd) - if result and result.returncode == 0 and result.stdout.strip(): - return result.stdout.strip() - return None - - -def get_diff_base(base_sha: str, head_sha: str = "HEAD", cwd: str = ".") -> str: - """Returns the base commit to diff against head_sha. - - Uses git merge-base when available so that changes that landed on the base - branch after the feature branch diverged are not falsely attributed to the PR. - - Args: - base_sha: Base git commit SHA. - head_sha: Head git commit SHA or ref. - cwd: Directory where the git command should be executed. - - Returns: - The merge-base SHA if available, otherwise base_sha. - """ - merge_base = get_merge_base(base_sha, head_sha, cwd=cwd) - return merge_base if merge_base else base_sha - - -def get_changed_files( - base_sha: str, head_sha: str = "HEAD", cwd: str = "." -) -> List[str]: - """Returns list of changed filepaths between base_sha and head_sha. - - Args: - base_sha: Base git commit SHA. - head_sha: Head git commit SHA or ref. - cwd: Directory where the git command should be executed. - - Returns: - List of relative paths of changed files. - """ - diff_base = get_diff_base(base_sha, head_sha, cwd=cwd) - command = ["git", "diff", "--name-only", diff_base, head_sha] - result = subprocess.run( - command, cwd=cwd, capture_output=True, text=True, check=True - ) - return [line.strip() for line in result.stdout.splitlines() if line.strip()] - - -def touches_build_keywords( - changed_files: Sequence[str], - base_sha: str, - head_sha: str, - cwd: str = ".", -) -> bool: - """Returns True if any BUILD or .bzl diff contains sensitive keywords.""" - build_or_bzl_files = [ - file_path - for file_path in changed_files - if file_path.endswith((".bzl", "BUILD", "BUILD.bazel")) - ] - if not build_or_bzl_files: - return False - - try: - diff_base = get_diff_base(base_sha, head_sha, cwd=cwd) - command = [ - "git", - "diff", - "-U0", - diff_base, - head_sha, - "--", - *build_or_bzl_files, - ] - result = subprocess.run( - command, cwd=cwd, capture_output=True, text=True, check=True - ) - diff_text = result.stdout - for keyword in FULL_RUN_BUILD_KEYWORDS: - if keyword in diff_text: - logging.info( - "Found build keyword '%s' in diff; forcing full run", keyword - ) - return True - except (OSError, subprocess.SubprocessError) as e: - logging.warning("Failed to check build keywords: %s", e) - return True - return False - - def normalize_target_label(label: str) -> str: """Normalizes Bazel target labels to standard repo forms.""" label = label.strip() @@ -537,6 +377,59 @@ def get_cquery_command_options(build: BuildConfig) -> List[str]: return options +# b/519689071: keywords in BUILD/.bzl changes that signal unconfigured hashing +# under-selection (label_flag, config_setting, alias). +FULL_RUN_BUILD_KEYWORDS: Tuple[str, ...] = ( + "label_flag(", + "config_setting(", + "alias(", +) + + +def _touches_build_keywords( + changed_files: Sequence[str], + base_sha: str, + head_sha: str, + cwd: str = ".", +) -> bool: + """Returns True if any BUILD or .bzl diff contains sensitive keywords.""" + build_or_bzl_files = [ + file_path + for file_path in changed_files + if file_path.endswith((".bzl", "BUILD", "BUILD.bazel")) + ] + if not build_or_bzl_files: + return False + + try: + diff_base = change_detector.get_diff_base(base_sha, head_sha, cwd=cwd) + safe_head = change_detector.validate_git_ref(head_sha) + command = [ + "git", + "diff", + "-U0", + "--end-of-options", + diff_base, + safe_head, + "--", + *build_or_bzl_files, + ] + result = subprocess.run( + command, cwd=cwd, capture_output=True, text=True, check=True + ) + diff_text = result.stdout + for keyword in FULL_RUN_BUILD_KEYWORDS: + if keyword in diff_text: + logging.info( + "Found build keyword '%s' in diff; forcing full run", keyword + ) + return True + except (OSError, subprocess.SubprocessError, ValueError) as e: + logging.warning("Failed to check build keywords: %s", e) + return True + return False + + def compute_impacted_targets( build: BuildConfig, base_sha: str, @@ -564,29 +457,38 @@ def compute_impacted_targets( if dest_jar_path is None: dest_jar_path = os.path.join(tempfile.gettempdir(), "bazel-diff.jar") + try: + base_sha = change_detector.validate_git_ref(base_sha) + head_sha = change_detector.validate_git_ref(head_sha or "HEAD") + except ValueError as e: + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason=f"Invalid git revision parameter: {e}", + elapsed_seconds=time.time() - start_time, + ) + # Record original ref currently checked out so we can faithfully restore it original_head = head_sha - result = _run_git_best_effort(["rev-parse", "HEAD"], cwd=workspace_dir) + result = change_detector.run_git_best_effort( + ["rev-parse", "HEAD"], cwd=workspace_dir + ) if result and result.returncode == 0 and result.stdout.strip(): - original_head = result.stdout.strip() + try: + original_head = change_detector.validate_git_ref(result.stdout.strip()) + except ValueError: + original_head = head_sha - if not head_sha or head_sha == "HEAD": + if head_sha == "HEAD": head_sha = original_head - # Fetch base SHA if shallow - _run_git_best_effort( - ["fetch", "--depth=1", "origin", base_sha], cwd=workspace_dir - ) - - if not get_merge_base(base_sha, head_sha, cwd=workspace_dir): - _run_git_best_effort(["fetch", "--unshallow", "origin"], cwd=workspace_dir) + change_detector.ensure_base_fetched(base_sha, head_sha, cwd=workspace_dir) # 1. Changed files analysis try: - changed_files = get_changed_files( + changed_files = change_detector.get_changed_files( base_sha, head_sha, cwd=workspace_dir ) - except (OSError, subprocess.SubprocessError) as e: + except (OSError, subprocess.SubprocessError, ValueError) as e: return BazelDiffDecision( decision=BazelDiffDecisionType.FULL, reason=f"Failed to get git changed files: {e}", @@ -594,7 +496,7 @@ def compute_impacted_targets( ) changed_count = len(changed_files) - if is_global_config_changed(changed_files): + if change_detector.is_global_config_changed(changed_files): return BazelDiffDecision( decision=BazelDiffDecisionType.FULL, reason="Global Bazel configuration modified", @@ -602,7 +504,7 @@ def compute_impacted_targets( elapsed_seconds=time.time() - start_time, ) - if touches_build_keywords( + if _touches_build_keywords( changed_files, base_sha, head_sha, cwd=workspace_dir ): return BazelDiffDecision( @@ -614,14 +516,6 @@ def compute_impacted_targets( elapsed_seconds=time.time() - start_time, ) - if is_docs_or_metadata_only(changed_files): - return BazelDiffDecision( - decision=BazelDiffDecisionType.SKIP, - reason="Only documentation or repository metadata modified", - changed_files_count=changed_count, - elapsed_seconds=time.time() - start_time, - ) - # 2. Download bazel-diff try: jar_path = download_and_verify_jar(dest_jar_path) @@ -645,11 +539,15 @@ def compute_impacted_targets( if build.startup_options: startup_options = dict_to_cli_options(build.startup_options) - diff_base = get_diff_base(base_sha, head_sha, cwd=workspace_dir) try: + diff_base = change_detector.validate_git_ref( + change_detector.get_diff_base(base_sha, head_sha, cwd=workspace_dir) + ) + safe_head = change_detector.validate_git_ref(head_sha) + # Generate base hashes subprocess.run( - ["git", "checkout", "-f", diff_base], + ["git", "checkout", "-f", "--end-of-options", diff_base], cwd=workspace_dir, check=True, capture_output=True, @@ -665,7 +563,7 @@ def compute_impacted_targets( # Generate head hashes subprocess.run( - ["git", "checkout", "-f", head_sha], + ["git", "checkout", "-f", "--end-of-options", safe_head], cwd=workspace_dir, check=True, capture_output=True, @@ -700,7 +598,16 @@ def compute_impacted_targets( ) finally: # Guaranteed restoration to original HEAD - _run_git_best_effort(["checkout", "-f", original_head], cwd=workspace_dir) + try: + safe_original_head = change_detector.validate_git_ref(original_head) + change_detector.run_git_best_effort( + ["checkout", "-f", "--end-of-options", safe_original_head], + cwd=workspace_dir, + ) + except ValueError: + logging.warning( + "Skipping git checkout restore for invalid ref: %r", original_head + ) target_patterns = build.target_patterns filtered_targets = filter_and_normalize_targets( diff --git a/third_party/xla/build_tools/ci/bazel_diff_test.py b/third_party/xla/build_tools/ci/bazel_diff_test.py index f663d36e0e943d..0bb6384ac60d37 100644 --- a/third_party/xla/build_tools/ci/bazel_diff_test.py +++ b/third_party/xla/build_tools/ci/bazel_diff_test.py @@ -16,6 +16,7 @@ import hashlib import os +import subprocess import tempfile from unittest import mock @@ -23,43 +24,11 @@ from absl.testing import parameterized from build_tools.ci import bazel_diff +from build_tools.ci import change_detector class BazelDiffTest(parameterized.TestCase): - @parameterized.parameters( - (["MODULE.bazel"], True), - (["REPO.bazel"], True), - (["WORKSPACE"], True), - (["WORKSPACE.bzlmod"], True), - ([".bazelrc"], True), - ([".bazelversion"], True), - (["xla/flags.bzl"], False), - (["third_party/tsl/tsl/platform/default/rules.bzl"], False), - (["third_party/tsl/REPO.bazel"], True), - (["third_party/tsl/tsl/BUILD"], False), - (["third_party/tsl/tsl/platform/denormal.cc"], False), - (["xla/service/cpu/cpu_compiler.cc"], False), - (["docs/overview.md", "xla/BUILD"], False), - ) - def test_is_global_config_changed(self, changed_files, expected): - self.assertEqual( - bazel_diff.is_global_config_changed(changed_files), expected - ) - - @parameterized.parameters( - (["README.md", "docs/index.md"], True), - (["OWNERS", "LICENSE"], True), - ([".gitignore", ".vscode/settings.json"], True), - (["logo.png", "diagram.svg"], True), - (["README.md", "xla/service/hlo_parser.cc"], False), - ([], False), - ) - def test_is_docs_or_metadata_only(self, changed_files, expected): - self.assertEqual( - bazel_diff.is_docs_or_metadata_only(changed_files), expected - ) - @parameterized.parameters( ("//xla/service:cpu_compiler", "//xla/service:cpu_compiler"), ("@@//xla/python:types", "//xla/python:types"), @@ -156,6 +125,61 @@ def test_target_allowed_by_patterns(self): ) ) + @mock.patch.object(bazel_diff.subprocess, "run", autospec=True, spec_set=True) + @mock.patch.object( + bazel_diff.change_detector, "get_diff_base", autospec=True, spec_set=True + ) + def test_touches_build_keywords_sensitive(self, mock_base, mock_run): + mock_base.return_value = "base_sha" + mock_run.return_value = subprocess.CompletedProcess( + args=["git", "diff"], + returncode=0, + stdout=( + "+ label_flag(name = 'custom_flag', default = '//xla/pkg:target')\n" + ), + ) + self.assertTrue( + bazel_diff._touches_build_keywords( + ["xla/flags.bzl"], "base_sha", "head_sha" + ) + ) + + @mock.patch.object(bazel_diff.subprocess, "run", autospec=True, spec_set=True) + @mock.patch.object( + bazel_diff.change_detector, "get_diff_base", autospec=True, spec_set=True + ) + def test_touches_build_keywords_benign_bzl(self, mock_base, mock_run): + mock_base.return_value = "base_sha" + mock_run.return_value = subprocess.CompletedProcess( + args=["git", "diff"], + returncode=0, + stdout=( + "+def my_macro(name, deps=[]):\n" + "+ native.cc_library(name='hlo_alias_analysis')\n" + ), + ) + self.assertFalse( + bazel_diff._touches_build_keywords( + ["xla/service/custom.bzl"], "base_sha", "head_sha" + ) + ) + mock_run.assert_called_once_with( + [ + "git", + "diff", + "-U0", + "--end-of-options", + "base_sha", + "head_sha", + "--", + "xla/service/custom.bzl", + ], + cwd=".", + capture_output=True, + text=True, + check=True, + ) + def test_filter_and_normalize_targets_with_patterns(self): raw_targets = [ "//xla/service/gpu:gpu_compiler", @@ -220,7 +244,9 @@ def test_report_decision(self): finally: os.unlink(summary_file) - @mock.patch.object(bazel_diff, "get_changed_files") + @mock.patch.object( + change_detector, "get_changed_files", autospec=True, spec_set=True + ) def test_compute_impacted_targets_global_config_fallback(self, mock_changed): mock_changed.return_value = ["MODULE.bazel", "xla/service/cpu_compiler.cc"] build_mock = mock.MagicMock() @@ -230,54 +256,6 @@ def test_compute_impacted_targets_global_config_fallback(self, mock_changed): self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL) self.assertIn("Global Bazel configuration", decision.reason) - @mock.patch.object(bazel_diff.subprocess, "run") - @mock.patch.object(bazel_diff, "get_diff_base") - def test_touches_build_keywords_sensitive(self, mock_base, mock_run): - mock_base.return_value = "base_sha" - mock_run.return_value = mock.MagicMock( - returncode=0, - stdout=( - "+ label_flag(name = 'custom_flag', default = '//xla/pkg:target')\n" - ), - ) - self.assertTrue( - bazel_diff.touches_build_keywords( - ["xla/flags.bzl"], "base_sha", "head_sha" - ) - ) - - @mock.patch.object(bazel_diff.subprocess, "run") - @mock.patch.object(bazel_diff, "get_diff_base") - def test_touches_build_keywords_benign_bzl(self, mock_base, mock_run): - mock_base.return_value = "base_sha" - mock_run.return_value = mock.MagicMock( - returncode=0, - stdout=( - "+def my_macro(name, deps=[]):\n" - "+ native.cc_library(name='hlo_alias_analysis')\n" - ), - ) - self.assertFalse( - bazel_diff.touches_build_keywords( - ["xla/service/custom.bzl"], "base_sha", "head_sha" - ) - ) - mock_run.assert_called_once_with( - [ - "git", - "diff", - "-U0", - "base_sha", - "head_sha", - "--", - "xla/service/custom.bzl", - ], - cwd=".", - capture_output=True, - text=True, - check=True, - ) - def test_get_cquery_command_options(self): build_mock = mock.MagicMock() build_mock.configs = ( @@ -328,8 +306,12 @@ def test_run_generate_hashes_cquery_options(self, mock_run): args, ) - @mock.patch.object(bazel_diff, "touches_build_keywords") - @mock.patch.object(bazel_diff, "get_changed_files") + @mock.patch.object( + bazel_diff, "_touches_build_keywords", autospec=True, spec_set=True + ) + @mock.patch.object( + change_detector, "get_changed_files", autospec=True, spec_set=True + ) def test_compute_impacted_targets_build_keywords_fallback( self, mock_changed, mock_keywords ): @@ -342,17 +324,9 @@ def test_compute_impacted_targets_build_keywords_fallback( self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL) self.assertIn("label_flag, config_setting, or alias", decision.reason) - @mock.patch.object(bazel_diff, "get_changed_files") - def test_compute_impacted_targets_docs_only_skip(self, mock_changed): - mock_changed.return_value = ["README.md", "docs/architecture.md"] - build_mock = mock.MagicMock() - decision = bazel_diff.compute_impacted_targets( - build_mock, "base_sha", "head_sha" - ) - self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.SKIP) - self.assertIn("Only documentation", decision.reason) - - @mock.patch.object(bazel_diff, "get_changed_files") + @mock.patch.object( + change_detector, "get_changed_files", autospec=True, spec_set=True + ) def test_compute_impacted_targets_git_fail_open(self, mock_changed): mock_changed.side_effect = OSError("git failure") build_mock = mock.MagicMock() @@ -362,30 +336,6 @@ def test_compute_impacted_targets_git_fail_open(self, mock_changed): self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL) self.assertIn("Failed to get git changed files", decision.reason) - @mock.patch.object(bazel_diff.subprocess, "run") - def test_get_merge_base_success(self, mock_run): - mock_run.return_value = mock.MagicMock(returncode=0, stdout="base_commit\n") - self.assertEqual(bazel_diff.get_merge_base("base", "head"), "base_commit") - - @mock.patch.object(bazel_diff.subprocess, "run") - def test_get_merge_base_failure(self, mock_run): - mock_run.return_value = mock.MagicMock(returncode=128, stdout="") - self.assertIsNone(bazel_diff.get_merge_base("base", "head")) - - @mock.patch.object(bazel_diff, "get_merge_base") - def test_get_diff_base_with_merge_base(self, mock_mb): - mock_mb.return_value = "merge_base_sha" - self.assertEqual( - bazel_diff.get_diff_base("base_sha", "head_sha"), "merge_base_sha" - ) - - @mock.patch.object(bazel_diff, "get_merge_base") - def test_get_diff_base_fallback(self, mock_mb): - mock_mb.return_value = None - self.assertEqual( - bazel_diff.get_diff_base("base_sha", "head_sha"), "base_sha" - ) - if __name__ == "__main__": absltest.main() diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py index 47f3b412843569..8c4583fbcd9139 100755 --- a/third_party/xla/build_tools/ci/build.py +++ b/third_party/xla/build_tools/ci/build.py @@ -33,6 +33,7 @@ from typing import Any, ClassVar, Dict, List, Optional, Tuple from build_tools.ci import bazel_diff +from build_tools.ci import change_detector # TODO(ddunleavy): move this to the bazelrc _DEFAULT_BAZEL_OPTIONS = dict( @@ -960,6 +961,16 @@ def dump_all_build_commands(): sys.stdout.write(f"# END {build.type_}\n") +def get_xla_dir(build: Build) -> str: + """Resolves the local XLA git repository directory for the given build.""" + for mapping in (build.override_module, build.override_repository): + if "xla" in mapping: + expanded = os.path.expandvars(mapping["xla"]) + if not expanded.startswith("$"): + return expanded + return "." + + def _parse_args(): """Defines flags and parses args.""" parser = argparse.ArgumentParser(allow_abbrev=False) @@ -990,8 +1001,19 @@ def main(): build = Build.all_builds()[args.build] github_event = os.environ.get("GITHUB_EVENT_NAME", "").strip() is_presubmit = github_event in ("", "pull_request", "pull_request_target") - base_sha = os.environ.get("XLA_CI_BAZEL_DIFF_BASE_SHA", "").strip() - head_sha = os.environ.get("XLA_CI_BAZEL_DIFF_HEAD_SHA", "HEAD").strip() + raw_base_sha = os.environ.get("XLA_CI_BAZEL_DIFF_BASE_SHA", "").strip() + raw_head_sha = ( + os.environ.get("XLA_CI_BAZEL_DIFF_HEAD_SHA", "HEAD").strip() or "HEAD" + ) + try: + base_sha = ( + change_detector.validate_git_ref(raw_base_sha) if raw_base_sha else "" + ) + head_sha = change_detector.validate_git_ref(raw_head_sha) + except ValueError as e: + logging.warning("Ignoring invalid git revision in environment: %s", e) + base_sha = "" + head_sha = "HEAD" disabled = os.environ.get("XLA_CI_DISABLE_BAZEL_DIFF", "false").lower() in ( "true", "1", @@ -999,20 +1021,38 @@ def main(): ) target_pattern_file = None - if build.use_bazel_diff and is_presubmit and base_sha and not disabled: - logging.info( - "Running bazel-diff analysis (base=%s, head=%s)", + if is_presubmit and base_sha and not disabled: + xla_dir = get_xla_dir(build) + skip_decision = change_detector.evaluate_skip( base_sha, head_sha, + cwd=xla_dir, + is_jax_build=(build.repo == "google/jax"), ) - decision = bazel_diff.compute_impacted_targets(build, base_sha, head_sha) - bazel_diff.report_decision(decision, str(build.type_)) - - if decision.decision == bazel_diff.BazelDiffDecisionType.SKIP: - logging.info("bazel-diff: 0 impacted targets, skipping Bazel commands!") + if skip_decision.should_skip: + change_detector.report_skip(skip_decision, str(build.type_)) + logging.info( + "change-detector: skipping Bazel commands (%s)", + skip_decision.reason, + ) return - elif decision.decision == bazel_diff.BazelDiffDecisionType.IMPACTED: - target_pattern_file = decision.impacted_targets_file + + if build.use_bazel_diff: + logging.info( + "Running bazel-diff analysis (base=%s, head=%s)", + base_sha, + head_sha, + ) + decision = bazel_diff.compute_impacted_targets( + build, base_sha, head_sha, workspace_dir=xla_dir + ) + bazel_diff.report_decision(decision, str(build.type_)) + + if decision.decision == bazel_diff.BazelDiffDecisionType.SKIP: + logging.info("bazel-diff: 0 impacted targets, skipping Bazel commands!") + return + elif decision.decision == bazel_diff.BazelDiffDecisionType.IMPACTED: + target_pattern_file = decision.impacted_targets_file for command in build.commands(target_pattern_file=target_pattern_file): result = sh(command, check=False) diff --git a/third_party/xla/build_tools/ci/change_detector.py b/third_party/xla/build_tools/ci/change_detector.py new file mode 100644 index 00000000000000..bcf66b9efb7bcd --- /dev/null +++ b/third_party/xla/build_tools/ci/change_detector.py @@ -0,0 +1,389 @@ +# 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. +# ============================================================================ +"""Git diff inspection and file-based fast-skip detection for XLA and JAX CI. + +Provides generic git diff utilities operating on a caller-supplied directory +and evaluates whether a CI build can be skipped based on the modified files. +""" + +from collections.abc import Sequence +import dataclasses +import logging +import os +import re +import subprocess +import sys +import time +from typing import List, Optional, Tuple + +# Global configuration files whose modification invalidates incremental diffs +# and necessitates a full test suite run. +GLOBAL_BAZEL_CONFIG_PATTERNS: Tuple[str, ...] = ( + r"(^|/)MODULE\.bazel$", + r"(^|/)REPO\.bazel$", + r"(^|/)WORKSPACE(\.bzlmod)?$", + r"\.bazelrc$", + r"\.bazelversion$", +) + +# Strict allowlist pattern for Git revision arguments (hexadecimal commit SHAs +# or alphanumeric ref names not starting with '-') to prevent flag injection. +SAFE_GIT_REF_PATTERN = re.compile(r"^[A-Za-z0-9_][A-Za-z0-9_./~^-]*$") + +# Regex matching files that only contain documentation or repository metadata. +DOCS_OR_METADATA_PATTERN = re.compile( + r"(\.(md|rst|png|jpg|jpeg|svg|webp|gif)$" + r"|(^|/)OWNERS$" + r"|(^|/)LICENSE(\.(txt|md))?$" + r"|(^|/)\.clang-(format|tidy)$" + r"|(^|/)\.clangd$" + r"|(^|/)\.gitignore$)" +) + +# Regex matching BUILD, .bzl, or .proto files that could alter build graphs. +BUILD_OR_PROTO_PATTERN = re.compile(r"((^|/)BUILD(\.bazel)?$|\.bzl$|\.proto$)") + +# Regex matching XLA unit test, benchmark, and HLO snapshot files. +XLA_TEST_OR_BENCHMARK_PATTERN = re.compile( + r"((_test|_tests|_test_gpu|_test_cpu|_benchmark)\.(cc|h|cu\.cc|py)$" + r"|\.hlo$)" +) + + +@dataclasses.dataclass(frozen=True) +class SkipDecision: + """Result of evaluating whether a CI build can be skipped. + + Attributes: + should_skip: Whether the CI build can be safely skipped. + reason: Human-readable explanation for the decision. + changed_files_count: Number of git changed files between base and head. + elapsed_seconds: Wall-clock duration of the skip evaluation in seconds. + """ + + should_skip: bool + reason: str = "" + changed_files_count: int = 0 + elapsed_seconds: float = 0.0 + + +def validate_git_ref(ref: str) -> str: + """Validates that ref is a safe Git revision string and cannot inject flags. + + Args: + ref: Git commit SHA or ref name to validate. + + Returns: + The validated ref string. + + Raises: + ValueError: If ref is empty, starts with '-', or contains disallowed chars. + """ + if not ref or ref.startswith("-") or not SAFE_GIT_REF_PATTERN.fullmatch(ref): + raise ValueError(f"Invalid or unsafe git revision parameter: {ref!r}") + return ref + + +def matches_any_pattern(path: str, patterns: Sequence[str]) -> bool: + """Returns True if path matches any regex pattern in patterns.""" + return any(re.search(pattern, path) for pattern in patterns) + + +def is_global_config_changed(changed_files: Sequence[str]) -> bool: + """Returns True if any changed file matches global bazel config patterns.""" + return any( + matches_any_pattern(file_path, GLOBAL_BAZEL_CONFIG_PATTERNS) + for file_path in changed_files + ) + + +def is_docs_or_metadata_only(changed_files: Sequence[str]) -> bool: + """Returns True if all changed files are docs or non-build metadata.""" + if not changed_files: + return False + return all( + bool(DOCS_OR_METADATA_PATTERN.search(file_path)) + for file_path in changed_files + ) + + +def run_git_best_effort( + args: Sequence[str], cwd: str = "." +) -> Optional[subprocess.CompletedProcess[str]]: + """Runs a git command best-effort, returning None on OS/subprocess error.""" + try: + return subprocess.run( + ["git", *args], + cwd=cwd, + capture_output=True, + text=True, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return None + + +def get_merge_base( + base_sha: str, head_sha: str = "HEAD", cwd: str = "." +) -> Optional[str]: + """Computes git merge-base between base_sha and head_sha. + + Args: + base_sha: Base git commit SHA. + head_sha: Head git commit SHA or ref. + cwd: Directory where the git command should be executed. + + Returns: + The merge base commit SHA as a string, or None if not found. + """ + try: + safe_base = validate_git_ref(base_sha) + safe_head = validate_git_ref(head_sha) + except ValueError: + return None + result = run_git_best_effort( + ["merge-base", "--", safe_base, safe_head], cwd=cwd + ) + if result and result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + return None + + +def get_diff_base(base_sha: str, head_sha: str = "HEAD", cwd: str = ".") -> str: + """Returns the base commit to diff against head_sha. + + Uses git merge-base when available so that changes that landed on the base + branch after the feature branch diverged are not falsely attributed to the PR. + + Args: + base_sha: Base git commit SHA. + head_sha: Head git commit SHA or ref. + cwd: Directory where the git command should be executed. + + Returns: + The merge-base SHA if available, otherwise validated base_sha. + """ + safe_base = validate_git_ref(base_sha) + safe_head = validate_git_ref(head_sha) + merge_base = get_merge_base(safe_base, safe_head, cwd=cwd) + return validate_git_ref(merge_base) if merge_base else safe_base + + +def ensure_base_fetched( + base_sha: str, head_sha: str = "HEAD", cwd: str = "." +) -> None: + """Best-effort fetches base_sha and unshallows if merge-base is missing.""" + try: + safe_base = validate_git_ref(base_sha) + safe_head = validate_git_ref(head_sha) + except ValueError: + return + run_git_best_effort( + ["fetch", "--depth=1", "--", "origin", safe_base], cwd=cwd + ) + if not get_merge_base(safe_base, safe_head, cwd=cwd): + run_git_best_effort(["fetch", "--unshallow", "origin"], cwd=cwd) + + +def get_changed_files( + base_sha: str, head_sha: str = "HEAD", cwd: str = "." +) -> List[str]: + """Returns list of changed filepaths between base_sha and head_sha. + + Args: + base_sha: Base git commit SHA. + head_sha: Head git commit SHA or ref. + cwd: Directory where the git command should be executed. + + Returns: + List of relative paths of changed files. + """ + diff_base = get_diff_base(base_sha, head_sha, cwd=cwd) + safe_head = validate_git_ref(head_sha) + command = [ + "git", + "diff", + "--name-only", + "--end-of-options", + diff_base, + safe_head, + "--", + ] + result = subprocess.run( + command, cwd=cwd, capture_output=True, text=True, check=True + ) + return [ + line.strip().replace("\\", "/") + for line in result.stdout.splitlines() + if line.strip() + ] + + +def is_build_or_proto_or_ci_core_file(file_path: str) -> bool: + """Returns True if file_path requires a full run on JAX presubmits.""" + normalized = file_path.replace("\\", "/") + if matches_any_pattern(normalized, GLOBAL_BAZEL_CONFIG_PATTERNS): + return True + if BUILD_OR_PROTO_PATTERN.search(normalized): + return True + if normalized.startswith((".github/", "build_tools/")) and not ( + DOCS_OR_METADATA_PATTERN.search(normalized) + ): + return True + return False + + +def is_unimpacted_for_all_jax(file_path: str) -> bool: + """Returns True if file_path is a test-only or non-code file for JAX CI.""" + normalized = file_path.replace("\\", "/") + if is_build_or_proto_or_ci_core_file(normalized): + return False + if DOCS_OR_METADATA_PATTERN.search(normalized): + return True + if ( + normalized.startswith("xla/") + and not normalized.startswith("xla/mosaic/") + and XLA_TEST_OR_BENCHMARK_PATTERN.search(normalized) + ): + return True + return False + + +def evaluate_skip( + base_sha: str, + head_sha: str = "HEAD", + cwd: str = ".", + *, + is_jax_build: bool = False, +) -> SkipDecision: + """Evaluates whether a CI build can be skipped based on git changed files. + + Guaranteed fail-open: any git, OS, or validation failure returns + should_skip=False. + + Args: + base_sha: Base git commit SHA to diff against. + head_sha: Head git commit SHA or ref. + cwd: Path to the git repository root to inspect. + is_jax_build: Whether the build under evaluation is a JAX CI job. + + Returns: + A SkipDecision indicating whether the build can be skipped and why. + """ + start_time = time.time() + cwd = os.path.abspath(cwd) + + try: + base_sha = validate_git_ref(base_sha) + head_sha = validate_git_ref(head_sha or "HEAD") + except ValueError as e: + return SkipDecision( + should_skip=False, + reason=f"Invalid git revision parameter: {e}", + elapsed_seconds=time.time() - start_time, + ) + + ensure_base_fetched(base_sha, head_sha, cwd=cwd) + + try: + changed_files = get_changed_files(base_sha, head_sha, cwd=cwd) + except (OSError, subprocess.SubprocessError, ValueError) as e: + return SkipDecision( + should_skip=False, + reason=f"Failed to get git changed files: {e}", + elapsed_seconds=time.time() - start_time, + ) + + changed_count = len(changed_files) + if not changed_files: + return SkipDecision( + should_skip=False, + reason="No changed files detected", + changed_files_count=0, + elapsed_seconds=time.time() - start_time, + ) + + if is_docs_or_metadata_only(changed_files): + return SkipDecision( + should_skip=True, + reason="Only documentation or repository metadata modified", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + if not is_jax_build: + return SkipDecision( + should_skip=False, + reason="Non-documentation changes require target analysis", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + if any(is_build_or_proto_or_ci_core_file(f) for f in changed_files): + return SkipDecision( + should_skip=False, + reason="Build configuration, proto, or CI workflow modified", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + if all(is_unimpacted_for_all_jax(f) for f in changed_files): + return SkipDecision( + should_skip=True, + reason="Modified files do not affect JAX builds", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + return SkipDecision( + should_skip=False, + reason="Changed files may impact JAX build targets", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + +def report_skip(decision: SkipDecision, build_name: str) -> None: + """Emits skip decision summary to stdout and appends to $GITHUB_STEP_SUMMARY.""" + status = "SKIP" if decision.should_skip else "RUN" + summary_line = ( + f"change-detector Build='{build_name}': " + f"Decision={status}, " + f"ChangedFiles={decision.changed_files_count}, " + f"Reason='{decision.reason}', " + f"Duration={decision.elapsed_seconds:.2f}s" + ) + separator = "=" * 55 + sys.stdout.write(f"\n{separator}\n") + sys.stdout.write(f"{summary_line}\n") + sys.stdout.write(f"{separator}\n\n") + sys.stdout.flush() + + step_summary_file = os.environ.get("GITHUB_STEP_SUMMARY") + if step_summary_file: + try: + with open(step_summary_file, "a", encoding="utf-8") as f: + f.write("### Change Detector Analysis\n") + f.write( + "| Parameter | Value |\n" + "|---|---|\n" + f"| **Build** | `{build_name}` |\n" + f"| **Decision** | **`{status}`** |\n" + f"| **Changed files** | {decision.changed_files_count} |\n" + f"| **Reason** | {decision.reason} |\n" + f"| **Duration** | {decision.elapsed_seconds:.2f}s |\n\n" + ) + except OSError as e: + logging.warning("Failed to write to GITHUB_STEP_SUMMARY: %s", e) diff --git a/third_party/xla/build_tools/ci/change_detector_test.py b/third_party/xla/build_tools/ci/change_detector_test.py new file mode 100644 index 00000000000000..ded1955bf7817b --- /dev/null +++ b/third_party/xla/build_tools/ci/change_detector_test.py @@ -0,0 +1,293 @@ +# 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. +# ============================================================================ +"""Tests for change_detector module.""" + +import os +import subprocess +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized + +from build_tools.ci import change_detector + + +class ChangeDetectorTest(parameterized.TestCase): + + @parameterized.parameters( + (["MODULE.bazel"], True), + (["REPO.bazel"], True), + (["WORKSPACE"], True), + (["WORKSPACE.bzlmod"], True), + ([".bazelrc"], True), + ([".bazelversion"], True), + (["xla/flags.bzl"], False), + (["third_party/tsl/tsl/platform/default/rules.bzl"], False), + (["third_party/tsl/REPO.bazel"], True), + (["third_party/tsl/tsl/BUILD"], False), + (["third_party/tsl/tsl/platform/denormal.cc"], False), + (["xla/service/cpu/cpu_compiler.cc"], False), + (["docs/overview.md", "xla/BUILD"], False), + ) + def test_is_global_config_changed(self, changed_files, expected): + self.assertEqual( + change_detector.is_global_config_changed(changed_files), expected + ) + + @parameterized.parameters( + (["README.md", "docs/index.md"], True), + (["OWNERS", "LICENSE", "LICENSE.txt"], True), + ([".gitignore", ".clang-format", ".clang-tidy", ".clangd"], True), + (["logo.png", "diagram.svg"], True), + (["README.md", "xla/service/hlo_parser.cc"], False), + (["LICENSE.py"], False), + (["LICENSE.sh"], False), + ([".clang_hook.sh"], False), + ([".gitignore.py"], False), + (["docs/script.py"], False), + ([], False), + ) + def test_is_docs_or_metadata_only(self, changed_files, expected): + self.assertEqual( + change_detector.is_docs_or_metadata_only(changed_files), expected + ) + + @parameterized.parameters( + ("26672b856b96dfb0a893f160ee61525526366663",), + ("HEAD",), + ("HEAD~1",), + ("origin/main",), + ("base_sha",), + ) + def test_validate_git_ref_valid(self, ref): + self.assertEqual(change_detector.validate_git_ref(ref), ref) + + @parameterized.parameters( + ("",), + ("-f",), + ("--upload-pack=evil",), + ("--output=/tmp/pwn",), + ("HEAD; rm -rf /",), + ("ref with spaces",), + ) + def test_validate_git_ref_invalid(self, ref): + with self.assertRaises(ValueError): + change_detector.validate_git_ref(ref) + + @mock.patch.object( + change_detector.subprocess, "run", autospec=True, spec_set=True + ) + def test_get_merge_base_success(self, mock_run): + mock_run.return_value = subprocess.CompletedProcess( + args=["git", "merge-base"], returncode=0, stdout="base_commit\n" + ) + self.assertEqual( + change_detector.get_merge_base("base", "head"), "base_commit" + ) + mock_run.assert_called_once_with( + ["git", "merge-base", "--", "base", "head"], + cwd=".", + capture_output=True, + text=True, + check=False, + ) + + @mock.patch.object( + change_detector.subprocess, "run", autospec=True, spec_set=True + ) + def test_get_merge_base_failure(self, mock_run): + mock_run.return_value = subprocess.CompletedProcess( + args=["git", "merge-base"], returncode=128, stdout="" + ) + self.assertIsNone(change_detector.get_merge_base("base", "head")) + + @mock.patch.object( + change_detector.subprocess, "run", autospec=True, spec_set=True + ) + def test_get_merge_base_rejects_flag_injection(self, mock_run): + self.assertIsNone( + change_detector.get_merge_base("--upload-pack=evil", "HEAD") + ) + mock_run.assert_not_called() + + @mock.patch.object( + change_detector, "get_merge_base", autospec=True, spec_set=True + ) + def test_get_diff_base_with_merge_base(self, mock_mb): + mock_mb.return_value = "merge_base_sha" + self.assertEqual( + change_detector.get_diff_base("base_sha", "head_sha"), "merge_base_sha" + ) + + @mock.patch.object( + change_detector, "get_merge_base", autospec=True, spec_set=True + ) + def test_get_diff_base_fallback(self, mock_mb): + mock_mb.return_value = None + self.assertEqual( + change_detector.get_diff_base("base_sha", "head_sha"), "base_sha" + ) + + @parameterized.named_parameters( + dict( + testcase_name="docs_only_xla_build", + changed_files=["README.md", "docs/architecture.md"], + is_jax_build=False, + expected_skip=True, + ), + dict( + testcase_name="docs_only_jax_build", + changed_files=["README.md", "OWNERS"], + is_jax_build=True, + expected_skip=True, + ), + dict( + testcase_name="non_docs_xla_build_defers_to_bazel_diff", + changed_files=["xla/service/gpu/fusion_merger_test.cc"], + is_jax_build=False, + expected_skip=False, + ), + dict( + testcase_name="xla_tests_and_benchmarks_skip_jax", + changed_files=[ + "xla/service/gpu/fusion_merger_test.cc", + "xla/hlo/ir/hlo_instruction_benchmark.cc", + "xla/service/gpu/tests/gemm_rewrite.hlo", + ], + is_jax_build=True, + expected_skip=True, + ), + dict( + testcase_name="mosaic_test_does_not_skip_jax", + changed_files=["xla/mosaic/gpu/mosaic_gpu_test.py"], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="workflow_and_build_tools_do_not_skip_jax", + changed_files=[ + ".github/workflows/benchmark_presubmit.yml", + "build_tools/rocm/run_xla.sh", + "build_tools/lint/diff_parser.py", + ], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="ci_yml_does_not_skip_jax", + changed_files=[".github/workflows/ci.yml"], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="build_tools_ci_does_not_skip_jax", + changed_files=["build_tools/ci/build.py"], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="gpu_sources_do_not_skip_jax", + changed_files=["xla/service/gpu/nvptx_compiler.cc"], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="build_file_does_not_skip_jax", + changed_files=[ + "xla/service/gpu/fusion_merger_test.cc", + "xla/service/gpu/BUILD", + ], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="proto_file_does_not_skip_jax", + changed_files=["xla/service/gpu/model/hlo_op_profile.proto"], + is_jax_build=True, + expected_skip=False, + ), + dict( + testcase_name="core_hlo_change_does_not_skip_jax", + changed_files=["xla/hlo/ir/hlo_instruction.cc"], + is_jax_build=True, + expected_skip=False, + ), + ) + @mock.patch.object( + change_detector, "ensure_base_fetched", autospec=True, spec_set=True + ) + @mock.patch.object( + change_detector, "get_changed_files", autospec=True, spec_set=True + ) + def test_evaluate_skip( + self, + mock_changed, + mock_fetch, + changed_files, + is_jax_build, + expected_skip, + ): + mock_changed.return_value = changed_files + decision = change_detector.evaluate_skip( + "base_sha", + "head_sha", + cwd="/workspace/openxla/xla", + is_jax_build=is_jax_build, + ) + self.assertEqual(decision.should_skip, expected_skip) + self.assertLen(changed_files, decision.changed_files_count) + mock_fetch.assert_called_once_with( + "base_sha", "head_sha", cwd="/workspace/openxla/xla" + ) + + @mock.patch.object( + change_detector, "ensure_base_fetched", autospec=True, spec_set=True + ) + @mock.patch.object( + change_detector, "get_changed_files", autospec=True, spec_set=True + ) + def test_evaluate_skip_git_fail_open(self, mock_changed, mock_fetch): + mock_changed.side_effect = OSError("git failure") + decision = change_detector.evaluate_skip( + "base_sha", "head_sha", is_jax_build=True + ) + self.assertFalse(decision.should_skip) + self.assertIn("Failed to get git changed files", decision.reason) + mock_fetch.assert_called_once() + + def test_report_skip(self): + summary_file = self.create_tempfile("step_summary.md") + with mock.patch.dict( + os.environ, {"GITHUB_STEP_SUMMARY": summary_file.full_path} + ): + decision = change_detector.SkipDecision( + should_skip=True, + reason="Modified files do not affect JAX builds", + changed_files_count=3, + elapsed_seconds=0.04, + ) + change_detector.report_skip( + decision, "BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS" + ) + + content = summary_file.read_text() + self.assertIn("Change Detector Analysis", content) + self.assertIn("JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS", content) + self.assertIn("SKIP", content) + self.assertIn("Modified files do not affect JAX builds", content) + + +if __name__ == "__main__": + absltest.main() diff --git a/third_party/xla/third_party/llvm/generated.patch b/third_party/xla/third_party/llvm/generated.patch index 634ed2ef4a9817..9da29843ede0af 100644 --- a/third_party/xla/third_party/llvm/generated.patch +++ b/third_party/xla/third_party/llvm/generated.patch @@ -1,664 +1,59 @@ Auto generated patch. Do not edit or delete it, even if empty. -diff -ruN --strip-trailing-cr a/clang/include/clang/AST/TypeBase.h b/clang/include/clang/AST/TypeBase.h ---- a/clang/include/clang/AST/TypeBase.h -+++ b/clang/include/clang/AST/TypeBase.h -@@ -2270,7 +2270,7 @@ - /// increments towards the beginning. - /// Positive non-zero number represents the index + 1. - /// Zero means this is not substituted from an expansion. -- unsigned PackIndex : 15; -+ unsigned PackIndex : 16; - }; - - class SubstPackTypeBitfields { -diff -ruN --strip-trailing-cr a/clang/test/SemaCXX/make_integer_seq.cpp b/clang/test/SemaCXX/make_integer_seq.cpp ---- a/clang/test/SemaCXX/make_integer_seq.cpp -+++ b/clang/test/SemaCXX/make_integer_seq.cpp -@@ -50,3 +50,8 @@ - - __make_integer_seq<__make_integer_seq, int, 10> PR28494; // expected-note{{different template parameters}} - // expected-error@make_integer_seq.cpp:* {{template argument for template template parameter must be a class template or type alias template}} -+ -+// The largest type pack whose pack indices fit in SubstTemplateTypeParmType. -+template using Expand = void(Ts...); -+template struct ExpandSeq { using type = Expand; }; -+using LargePack = __make_integer_seq::type; -diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp ---- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp -+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp -@@ -13901,47 +13901,12 @@ - return SDValue(); +diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h +--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h ++++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h +@@ -837,7 +837,7 @@ + + SDValue ScalarizeVecRes_ADDRSPACECAST(SDNode *N); + SDValue ScalarizeVecRes_BITCAST(SDNode *N); +- SDValue ScalarizeVecRes_BUILD_VECTOR(SDNode *N); ++ SDValue ScalarizeVecRes_BUILD_VECTOR_OR_SPLAT(SDNode *N); + SDValue ScalarizeVecRes_EXTRACT_SUBVECTOR(SDNode *N); + SDValue ScalarizeVecRes_FP_ROUND(SDNode *N); + SDValue ScalarizeVecRes_CONVERT_FROM_ARBITRARY_FP(SDNode *N); +diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp +--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp ++++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp +@@ -63,7 +63,10 @@ + break; + case ISD::MERGE_VALUES: R = ScalarizeVecRes_MERGE_VALUES(N, ResNo);break; + case ISD::BITCAST: R = ScalarizeVecRes_BITCAST(N); break; +- case ISD::BUILD_VECTOR: R = ScalarizeVecRes_BUILD_VECTOR(N); break; ++ case ISD::SPLAT_VECTOR: ++ case ISD::BUILD_VECTOR: ++ R = ScalarizeVecRes_BUILD_VECTOR_OR_SPLAT(N); ++ break; + case ISD::EXTRACT_SUBVECTOR: R = ScalarizeVecRes_EXTRACT_SUBVECTOR(N); break; + case ISD::FP_ROUND: R = ScalarizeVecRes_FP_ROUND(N); break; + case ISD::CONVERT_FROM_ARBITRARY_FP: +@@ -475,10 +478,10 @@ + NewVT, Op); } --/// Check if Mask defines a known constant set of enabled lanes, where only the --/// first N lanes are enabled. N is returned if so. --static uint64_t calculateConstantLowMaskLanes(SDValue Mask) { -- // We expect masks for masked load/store to be i1 predicates. -- if (Mask.getValueType().getScalarSizeInBits() != 1) -- return 0; -- -- if (Mask.getOpcode() == ISD::BUILD_VECTOR) { -- unsigned NumOnes = 0; -- auto Op = Mask->op_begin(); -- for (; Op != Mask->op_end(); ++Op) { -- if (!isOneConstant(*Op)) -- break; -- ++NumOnes; -- } -- -- if (NumOnes == 0) -- return 0; -- -- for (; Op != Mask->op_end(); ++Op) -- if (!isNullConstant(*Op)) -- return 0; -- return NumOnes; -- } -- -- if (Mask.getOpcode() == ISD::GET_ACTIVE_LANE_MASK && -- isNullConstant(Mask.getOperand(0)) && -- isa(Mask.getOperand(1)) && -- Mask.getOperand(1).getValueType().getSizeInBits() <= 64) -- return Mask.getConstantOperandVal(1); -- -- return 0; --} -- - SDValue DAGCombiner::visitMSTORE(SDNode *N) { - MaskedStoreSDNode *MST = cast(N); - SDValue Mask = MST->getMask(); - SDValue Chain = MST->getChain(); - SDValue Value = MST->getValue(); - SDValue Ptr = MST->getBasePtr(); -- EVT VT = Value.getValueType(); - - // Zap masked stores with a zero mask. - if (ISD::isConstantSplatVectorAllZeros(Mask.getNode())) -@@ -13964,46 +13929,21 @@ - } - } - -- // If this is a masked store with an all ones mask, we can use a unmasked -- // store. -+ // If this is a masked load with an all ones mask, we can use a unmasked load. - // FIXME: Can we do this for indexed, compressing, or truncating stores? -- if (MST->isUnindexed() && !MST->isCompressingStore() && -- !MST->isTruncatingStore()) { -- if (ISD::isConstantSplatVectorAllOnes(Mask.getNode())) -- return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), -- MST->getBasePtr(), MST->getPointerInfo(), -- MST->getBaseAlign(), MST->getMemOperand()->getFlags(), -- MST->getAAInfo()); -- -- // Convert a masked_store with constant getactivelanemask input mask to a -- // standard store. -- if (uint64_t Lanes = calculateConstantLowMaskLanes(Mask)) { -- if (Lanes < VT.getVectorMinNumElements() && isPowerOf2_32(Lanes)) { -- EVT SubVT = EVT::getVectorVT(*DAG.getContext(), -- VT.getVectorElementType(), Lanes); -- unsigned IsFast = 0; -- if ((!LegalTypes || isTypeLegal(SubVT)) && -- TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), -- SubVT, MST->getAddressSpace(), -- MST->getBaseAlign(), -- MST->getMemOperand()->getFlags(), &IsFast) && -- IsFast) { -- SDLoc DL(N); -- SDValue Ext = DAG.getExtractSubvector(DL, SubVT, MST->getValue(), 0); -- return DAG.getStore(MST->getChain(), DL, Ext, MST->getBasePtr(), -- MST->getPointerInfo(), MST->getBaseAlign(), -- MST->getMemOperand()->getFlags(), -- MST->getAAInfo()); -- } -- } -- } -- } -+ if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MST->isUnindexed() && -+ !MST->isCompressingStore() && !MST->isTruncatingStore()) -+ return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), -+ MST->getBasePtr(), MST->getPointerInfo(), -+ MST->getBaseAlign(), MST->getMemOperand()->getFlags(), -+ MST->getAAInfo()); - - // Try transforming N to an indexed store. - if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N)) - return SDValue(N, 0); - -- if (MST->isTruncatingStore() && MST->isUnindexed() && VT.isInteger() && -+ if (MST->isTruncatingStore() && MST->isUnindexed() && -+ Value.getValueType().isInteger() && - (!isa(Value) || - !cast(Value)->isOpaque())) { - APInt TruncDemandedBits = +-SDValue DAGTypeLegalizer::ScalarizeVecRes_BUILD_VECTOR(SDNode *N) { ++SDValue DAGTypeLegalizer::ScalarizeVecRes_BUILD_VECTOR_OR_SPLAT(SDNode *N) { + EVT EltVT = N->getValueType(0).getVectorElementType(); + SDValue InOp = N->getOperand(0); +- // The BUILD_VECTOR operands may be of wider element types and ++ // The BUILD_VECTOR / SPLAT operands may be of wider element types and + // we may need to truncate them back to the requested return type. + if (EltVT.isInteger()) + return DAG.getNode(ISD::TRUNCATE, SDLoc(N), EltVT, InOp); diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll --- a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll +++ b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll -@@ -139,7 +139,14 @@ - define void @store_low2_v4f64(ptr %p, <4 x double> %a) vscale_range(2, 2) { - ; CHECK-LABEL: store_low2_v4f64: - ; CHECK: // %bb.0: --; CHECK-NEXT: str q0, [x0] -+; CHECK-NEXT: // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1 -+; CHECK-NEXT: ptrue p0.d, vl2 -+; CHECK-NEXT: ptrue p1.h, vl2 -+; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1 -+; CHECK-NEXT: splice z0.d, p0, { z0.d, z1.d } -+; CHECK-NEXT: punpklo p0.h, p1.b -+; CHECK-NEXT: punpklo p0.h, p0.b -+; CHECK-NEXT: st1d { z0.d }, p0, [x0] - ; CHECK-NEXT: ret - %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) - tail call void @llvm.masked.store(<4 x double> %a, ptr align 1 %p, <4 x i1> %m) -@@ -149,7 +156,13 @@ - define void @store_low2_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { - ; CHECK-LABEL: store_low2_v4i32: - ; CHECK: // %bb.0: --; CHECK-NEXT: str d0, [x0] -+; CHECK-NEXT: ptrue p0.h, vl2 -+; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 -+; CHECK-NEXT: mov z1.h, p0/z, #-1 // =0xffffffffffffffff -+; CHECK-NEXT: ptrue p0.s -+; CHECK-NEXT: sshll v1.4s, v1.4h, #0 -+; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 -+; CHECK-NEXT: st1w { z0.s }, p1, [x0] - ; CHECK-NEXT: ret - %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) - tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> %m) -@@ -159,7 +172,12 @@ - define void @store_low2c_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { - ; CHECK-LABEL: store_low2c_v4i32: - ; CHECK: // %bb.0: --; CHECK-NEXT: str d0, [x0] -+; CHECK-NEXT: adrp x8, .LCPI12_0 -+; CHECK-NEXT: ptrue p0.s -+; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 -+; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI12_0] -+; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 -+; CHECK-NEXT: st1w { z0.s }, p1, [x0] - ; CHECK-NEXT: ret - tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> ) - ret void -@@ -168,7 +186,11 @@ - define void @store_low2_nxv4f64(ptr %p, %a) { - ; CHECK-LABEL: store_low2_nxv4f64: - ; CHECK: // %bb.0: --; CHECK-NEXT: str q0, [x0] -+; CHECK-NEXT: ptrue p0.s, vl2 -+; CHECK-NEXT: punpkhi p1.h, p0.b -+; CHECK-NEXT: punpklo p0.h, p0.b -+; CHECK-NEXT: st1d { z1.d }, p1, [x0, #1, mul vl] -+; CHECK-NEXT: st1d { z0.d }, p0, [x0] - ; CHECK-NEXT: ret - %m = tail call @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 2) - tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) -@@ -203,10 +225,13 @@ - define void @store_low2_v8i16(ptr %p, <8 x i16> %a) vscale_range(2, 2) { - ; CHECK-LABEL: store_low2_v8i16: - ; CHECK: // %bb.0: -+; CHECK-NEXT: ptrue p0.b, vl2 - ; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 --; CHECK-NEXT: ptrue p0.s, vl2 --; CHECK-NEXT: uunpklo z0.s, z0.h --; CHECK-NEXT: st1h { z0.s }, p0, [x0] -+; CHECK-NEXT: mov z1.b, p0/z, #-1 // =0xffffffffffffffff -+; CHECK-NEXT: ptrue p0.h, vl8 -+; CHECK-NEXT: sshll v1.8h, v1.8b, #0 -+; CHECK-NEXT: cmpne p1.h, p0/z, z1.h, #0 -+; CHECK-NEXT: st1h { z0.h }, p1, [x0] - ; CHECK-NEXT: ret - %m = tail call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 0, i32 2) - tail call void @llvm.masked.store(<8 x i16> %a, ptr align 1 %p, <8 x i1> %m) -@@ -216,9 +241,8 @@ - define void @store_low2_nxv8i16(ptr %p, %a) { - ; CHECK-LABEL: store_low2_nxv8i16: - ; CHECK: // %bb.0: --; CHECK-NEXT: mov z1.h, z0.h[1] --; CHECK-NEXT: str h0, [x0] --; CHECK-NEXT: str h1, [x0, #2] -+; CHECK-NEXT: ptrue p0.h, vl2 -+; CHECK-NEXT: st1h { z0.h }, p0, [x0] - ; CHECK-NEXT: ret - %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 2) - tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) -@@ -239,7 +263,8 @@ - define void @store_low4_nxv8i16(ptr %p, %a) { - ; CHECK-LABEL: store_low4_nxv8i16: - ; CHECK: // %bb.0: --; CHECK-NEXT: str d0, [x0] -+; CHECK-NEXT: ptrue p0.h, vl4 -+; CHECK-NEXT: st1h { z0.h }, p0, [x0] - ; CHECK-NEXT: ret - %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 4) +@@ -245,3 +245,14 @@ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) -diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll ---- a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll -+++ b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll -@@ -6,43 +6,40 @@ - ; CHECK-LABEL: f0: - ; CHECK: // %bb.0: // %b0 - ; CHECK-NEXT: { --; CHECK-NEXT: r0 = add(r1,r0) -+; CHECK-NEXT: [[RADD:[r0-9]+]] = add([[R1:[r0-9]+]],[[R0:[r0-9]+]]) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: allocframe(r29,#128):raw -+; CHECK-NEXT: [[RLCPI:[r0-9]+]] = ##.LCPI0_0 - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: r2 = add(r0,#136) -+; CHECK-NEXT: [[RNEG:[r0-9]+]] = #-1 - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: r29 = and(r29,#-128) -+; CHECK-NEXT: [[RADD136:[r0-9]+]] = add([[RADD]],#136) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: v1 = vmem(r0+#1) -+; CHECK-NEXT: [[V1:[v0-9]+]] = vmem([[RLCPI]]+#0) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: v0 = vmemu(r2+#0) -+; CHECK-NEXT: [[RADD128:[r0-9]+]] = add([[RADD]],#128) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: r4 = add(r29,#0) -+; CHECK-NEXT: [[Q0:[q0-9]+]] = vand([[V1]],[[RNEG]]) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: v0.w = vadd(v1.w,v0.w) -+; CHECK-NEXT: [[V2:[v0-9]+]] = vmemu([[RADD136]]+#0) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: vmem(r4+#0) = v0 -+; CHECK-NEXT: [[V3:[v0-9]+]] = vmem([[RADD]]+#1) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: r2 = memw(r4+#0) -+; CHECK-NEXT: [[V0:[v0-9]+]].w = vadd([[V3]].w,[[V2]].w) - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: r3 = memw(r4+#4) -+; CHECK-NEXT: if ([[Q0]]) vmem([[RADD128]]+#0) = [[V0]] - ; CHECK-NEXT: } - ; CHECK-NEXT: { --; CHECK-NEXT: memd(r0+#128) = r3:2 --; CHECK-NEXT: } --; CHECK-NEXT: { --; CHECK-NEXT: r31:30 = dealloc_return(r30):raw -+; CHECK-NEXT: jumpr r31 - ; CHECK-NEXT: } - b0: - %v0 = add i32 %a1, 128 -diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll ---- a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll -+++ b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll -@@ -29,12 +29,20 @@ - define void @kmovkr_2() { - ; AVX512-LABEL: kmovkr_2: - ; AVX512: # %bb.0: # %alloca_21 --; AVX512-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] -+; AVX512-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] -+; AVX512-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] -+; AVX512-NEXT: kmovw %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xf8,0x92,0xc8] -+; AVX512-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] -+; AVX512-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] - ; AVX512-NEXT: retq # encoding: [0xc3] - ; - ; AVX512BW-LABEL: kmovkr_2: - ; AVX512BW: # %bb.0: # %alloca_21 --; AVX512BW-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] -+; AVX512BW-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] -+; AVX512BW-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] -+; AVX512BW-NEXT: kmovd %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xfb,0x92,0xc8] -+; AVX512BW-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] -+; AVX512BW-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] - ; AVX512BW-NEXT: retq # encoding: [0xc3] - alloca_21: - call void @llvm.masked.store.v4f32.p0(<4 x float> zeroinitializer, ptr null, i32 1, <4 x i1> ) -diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll ---- a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll -+++ b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll -@@ -1,21 +0,0 @@ --; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py --; RUN: llc < %s -mtriple=x86_64-linux-gnu -mattr=+avx,+avx2,+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl | FileCheck %s -- --define fastcc void @render_event(<12 x double> %0) { --; CHECK-LABEL: render_event: --; CHECK: # %bb.0: --; CHECK-NEXT: vmovlhps {{.*#+}} xmm6 = xmm6[0],xmm7[0] --; CHECK-NEXT: vmovlhps {{.*#+}} xmm4 = xmm4[0],xmm5[0] --; CHECK-NEXT: vinsertf128 $1, %xmm6, %ymm4, %ymm4 --; CHECK-NEXT: vmovlhps {{.*#+}} xmm2 = xmm2[0],xmm3[0] --; CHECK-NEXT: vmovlhps {{.*#+}} xmm0 = xmm0[0],xmm1[0] --; CHECK-NEXT: vinsertf128 $1, %xmm2, %ymm0, %ymm0 --; CHECK-NEXT: vinsertf64x4 $1, %ymm4, %zmm0, %zmm0 --; CHECK-NEXT: vmovsd {{.*#+}} xmm1 = mem[0],zero --; CHECK-NEXT: vmovsd %xmm1, 64 --; CHECK-NEXT: vmovups %zmm0, 0 --; CHECK-NEXT: vzeroupper --; CHECK-NEXT: retq -- call void @llvm.masked.store.v12f64.p0(<12 x double> %0, ptr null, <12 x i1> ) -- ret void --} -diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr179489.ll b/llvm/test/CodeGen/X86/pr179489.ll ---- a/llvm/test/CodeGen/X86/pr179489.ll -+++ b/llvm/test/CodeGen/X86/pr179489.ll -@@ -9,12 +9,14 @@ - ; X86-LABEL: foo: - ; X86: # %bb.0: # %entry - ; X86-NEXT: movl {{[0-9]+}}(%esp), %eax --; X86-NEXT: vpextrb $1, %xmm0, (%eax) -+; X86-NEXT: vmovw %xmm0, %ecx -+; X86-NEXT: movb %ch, (%eax) - ; X86-NEXT: retl - ; - ; X64-LABEL: foo: - ; X64: # %bb.0: # %entry --; X64-NEXT: vpextrb $1, %xmm0, (%rdi) -+; X64-NEXT: vmovw %xmm0, %eax -+; X64-NEXT: movb %ah, (%rdi) - ; X64-NEXT: retq - entry: - %i3 = shufflevector <8 x i16> %arg, <8 x i16> , <8 x i32> -diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr46532.ll b/llvm/test/CodeGen/X86/pr46532.ll ---- a/llvm/test/CodeGen/X86/pr46532.ll -+++ b/llvm/test/CodeGen/X86/pr46532.ll -@@ -7,6 +7,8 @@ - ; CHECK-NEXT: movq (%rax), %rax - ; CHECK-NEXT: vxorps %xmm0, %xmm0, %xmm0 - ; CHECK-NEXT: vmovaps %xmm0, 32(%rax) -+; CHECK-NEXT: vmovsd {{.*#+}} xmm0 = [4294967295,4294967295,0,0] -+; CHECK-NEXT: vmaskmovps %ymm0, %ymm0, (%rax) - while.1.body.preheader: - %0 = load ptr, ptr undef, align 8, !invariant.load !0, !dereferenceable !1, !align !2 - %1 = getelementptr inbounds i8, ptr %0, i64 32 -diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td ---- a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td -+++ b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td -@@ -1,71 +0,0 @@ --// RUN: split-file %s %t --// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ --// RUN: | FileCheck %t/valid.td --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-bug.td 2>&1 \ --// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=user-sgpr-init16-bug --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-priority.td 2>&1 \ --// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=required-export-priority --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-timing.td 2>&1 \ --// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=gfx125x-lowest-rate-wmma --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-banks.td 2>&1 \ --// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=ldsbankcount32 --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-lds-size.td 2>&1 \ --// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=addressablelocalmemorysize65536 --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/capability.td 2>&1 \ --// RUN: | FileCheck %t/capability.td -DFILE=%t/capability.td --implicit-check-not="error:" -- --//--- common.td --include "AMDGPU.td" -- --def EmptyGPU : AMDGPUProcessorModel<"gfx999", NoSchedModel, [], [9, 9, 9]>; -- --//--- valid.td --include "common.td" --// Check the real generic targets, including gfx11's combined workarounds, --// gfx12-5's fixup and WMMA timing, and gfx9-4's LDS size and bank count. --// CHECK: GK_GFX9_4_GENERIC --// CHECK: GK_GFX11_GENERIC --// CHECK: GK_GFX12_5_GENERIC -- --//--- missing-bug.td --include "common.td" --defvar CheckedFeature = FeatureUserSGPRInit16Bug; --include "missing-generic.td" -- --//--- missing-priority.td --include "common.td" --defvar CheckedFeature = FeatureRequiredExportPriority; --include "missing-generic.td" -- --//--- missing-timing.td --include "common.td" --defvar CheckedFeature = FeatureGFX125xLowestRateWMMA; --include "missing-generic.td" -- --//--- missing-banks.td --include "common.td" --defvar CheckedFeature = FeatureLDSBankCount32; --include "missing-generic.td" -- --//--- missing-lds-size.td --include "common.td" --defvar CheckedFeature = FeatureAddressableLocalMemorySize65536; --include "missing-generic.td" -- --//--- missing-generic.td --// Classified backend-only features still need support from a covered GPU. --// Included files use the native path separator after the include directory. --// CHECK: [[DIR]]{{[/\\]}}missing-generic.td:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature '[[FEATURE]]' not supported by any covered GPU --def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [CheckedFeature], [9, 9, 0]> { -- let CoveredGPUs = [EmptyGPU]; --} -- --//--- capability.td --include "common.td" --def WMMAGPU : AMDGPUProcessorModel<"gfx998", NoSchedModel, -- [FeatureWMMA256bInsts], [9, 9, 8]>; --// Instruction support keeps the default all-covered-GPUs rule. --// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature 'wmma-256b-insts' not supported by covered GPU 'gfx999' --def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [FeatureWMMA256bInsts], [9, 9, 0]> { -- let CoveredGPUs = [WMMAGPU, EmptyGPU]; --} -diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td ---- a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td -+++ b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td -@@ -1,33 +1,66 @@ - // RUN: split-file %s %t --// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ -+// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/valid.td \ - // RUN: | FileCheck %t/valid.td --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ - // RUN: | FileCheck %t/missing-any.td -DFILE=%t/missing-any.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ - // RUN: | FileCheck %t/missing-all.td -DFILE=%t/missing-all.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ - // RUN: | FileCheck %t/implied-all.td -DFILE=%t/implied-all.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ - // RUN: | FileCheck %t/missing-private.td -DFILE=%t/missing-private.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-private.td 2>&1 \ -+// RUN: | FileCheck %t/implied-private.td -DFILE=%t/implied-private.td --implicit-check-not="error:" -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ - // RUN: | FileCheck %t/numeric-mismatch.td -DFILE=%t/numeric-mismatch.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ - // RUN: | FileCheck %t/empty-members.td -DFILE=%t/empty-members.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ - // RUN: | FileCheck %t/self-member.td -DFILE=%t/self-member.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ - // RUN: | FileCheck %t/generic-member.td -DFILE=%t/generic-member.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ - // RUN: | FileCheck %t/pseudo-member.td -DFILE=%t/pseudo-member.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ - // RUN: | FileCheck %t/unregistered-member.td -DFILE=%t/unregistered-member.td --implicit-check-not="error:" --// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ -+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ - // RUN: | FileCheck %t/r600-member.td -DFILE=%t/r600-member.td --implicit-check-not="error:" - - //--- common.td --include "llvm/Target/Target.td" -+// Exercise the validation rules with small synthetic targets. The real targets -+// are validated when AMDGPUTargetParserDef.inc is generated during the build. -+// Stand-ins for the Target.td classes used by AMDGPUTargetParser.td, limited to -+// the fields the backend reads. -+class InlineBehavior; -+def InlineIgnore : InlineBehavior; -+def InlineDefault : InlineBehavior; -+class SubtargetFeature i = [], -+ InlineBehavior ib = InlineDefault> { -+ string Name = n; -+ string FieldName = f; -+ string Value = v; -+ string Desc = d; -+ list Implies = i; -+ InlineBehavior InlineBehavior = ib; -+} -+ -+class SchedMachineModel; -+def NoSchedModel : SchedMachineModel; -+class ProcessorItineraries; -+def NoItineraries : ProcessorItineraries; -+ -+class Processor f> { -+ string Name = n; -+ ProcessorItineraries ProcItin = pi; -+ list Features = f; -+} -+class ProcessorModel f> -+ : Processor { -+ SchedMachineModel SchedModel = m; -+} -+ - include "AMDGPUTargetParser.td" --def MyTarget : Target; - - def FeatureAll : SubtargetFeature<"all", "HasAll", "true", "">; - def FeatureAnyA : AMDGPUGenericAnyFeature<"any-a", "HasAnyA", "true", "">; -@@ -37,7 +70,13 @@ - def FeatureAnySmall : NumericFeature<4>; - def FeatureAnyB : NumericFeature<8>; - def FeatureAnyLarge : NumericFeature<16>; --def FeaturePrivate : AMDGPUGenericAnyFeature<"private", "HasPrivate", "true", "">; -+// Classification inherited through a subclass also applies to backend-only -+// features, which must participate in validation without being emitted. -+class PrivateFeature : AMDGPUGenericAnyFeature< -+ "private", "HasPrivate", "true", "">; -+def FeaturePrivate : PrivateFeature; -+def FeatureImpliesPrivate : SubtargetFeature< -+ "implies-private", "HasImpliesPrivate", "true", "", [FeaturePrivate]>; - def FeatureImpliesA : SubtargetFeature< - "implies-a", "HasImpliesA", "true", "", [FeatureAnyA]>; - def FeatureImpliesB : SubtargetFeature< -@@ -55,7 +94,7 @@ - def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, - [FeatureAll, FeatureAnyA, FeatureAnySmall], [9, 0, 0]>; - def GFX901 : AMDGPUProcessorModel<"gfx901", NoSchedModel, -- [FeatureAll, FeatureImpliesB, FeaturePrivate], [9, 0, 1]>; -+ [FeatureAll, FeatureImpliesB, FeatureImpliesPrivate], [9, 0, 1]>; - // Different members supply different marked features. Follow Implies on both - // the generic and the members, and check marked backend-only features too. - def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, -@@ -103,6 +142,15 @@ - let CoveredGPUs = [GFX900]; + ret void } - -+//--- implied-private.td -+include "common.td" -+def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [], [9, 0, 0]>; -+// An unclassified backend-only parent must not hide a classified dependency. -+// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-generic' exposes feature 'private' not supported by any covered GPU -+def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, [FeatureImpliesPrivate], [9, 0, 0]> { -+ let CoveredGPUs = [GFX900]; -+} + - //--- numeric-mismatch.td - include "common.td" - def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [FeatureAnySmall], [9, 0, 0]>; -diff -ruN --strip-trailing-cr a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp ---- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp -+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp -@@ -14,11 +14,12 @@ - using namespace mlir; - using namespace mlir::tosa; - --// Building this ~5,000-line generated initializer map under MemorySanitizer --// leads to extreme compile times (>3.5 minutes) during InstCombine and greedy --// register allocation due to shadow/origin instrumentation overhead. Disable --// optimization under MSan to keep build times manageable. --#if LLVM_MEMORY_SANITIZER_BUILD -+// Building this ~5,000-line generated initializer map under HWAddressSanitizer -+// or MemorySanitizer leads to extreme compile times (>3.5 minutes) during -+// InstCombine and greedy register allocation due to instrumentation -+// overhead. Disable optimization under MSan/HWAsan to keep build times -+// manageable. -+#if LLVM_MEMORY_SANITIZER_BUILD || LLVM_HWADDRESS_SANITIZER_BUILD - __attribute__((optnone)) - #else - LLVM_ATTRIBUTE_MINSIZE -diff -ruN --strip-trailing-cr a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel ---- a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel -+++ b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel -@@ -66,7 +66,7 @@ - - bool_flag( - name = "llvm_enable_lzma", -- build_setting_default = False, -+ build_setting_default = True, - ) - - config_setting( -@@ -74,11 +74,7 @@ - flag_values = {":llvm_enable_lzma": "true"}, - ) - --# Unlike zlib and zstd there is no liblzma module in the Bazel registry, so this --# links the system library rather than building one from source. That is why it --# is a plain cc_library with linkopts instead of a cc_library_wrapper, and why --# it defaults to off. --cc_library( -+cc_library_wrapper( - name = "lzma", - defines = select({ - ":llvm_lzma_enabled": [ -@@ -88,9 +84,9 @@ - "LLVM_ENABLE_LZMA=0", - ], - }), -- linkopts = select({ -+ deps = select({ - ":llvm_lzma_enabled": [ -- "-llzma", -+ "@llvm_xz//:lzma", - ], - "//conditions:default": [], - }), -diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel b/utils/bazel/MODULE.bazel ---- a/utils/bazel/MODULE.bazel -+++ b/utils/bazel/MODULE.bazel -@@ -29,6 +29,7 @@ - bazel_dep(name = "nanobind", version = "2.12.0.bcr.1") - bazel_dep(name = "libpfm", version = "4.13.0", repo_name = "pfm") - bazel_dep(name = "vulkan_headers", version = "1.4.349") -+bazel_dep(name = "xz", version = "5.4.5.bcr.9", repo_name = "llvm_xz") - - bazel_dep(name = "emsdk", version = "6.0.2", dev_dependency = True) - bazel_dep(name = "llvm", version = "0.8.20", dev_dependency = True) -diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel.lock b/utils/bazel/MODULE.bazel.lock ---- a/utils/bazel/MODULE.bazel.lock -+++ b/utils/bazel/MODULE.bazel.lock -@@ -55,7 +55,8 @@ - "https://bcr.bazel.build/modules/bazel_lib/3.0.0/MODULE.bazel": "22b70b80ac89ad3f3772526cd9feee2fa412c2b01933fea7ed13238a448d370d", - "https://bcr.bazel.build/modules/bazel_lib/3.1.0/MODULE.bazel": "6809765c14e3c766a9b9286c7b0ec56ed87a73326e48fe01749f0c0fdcfe3287", - "https://bcr.bazel.build/modules/bazel_lib/3.2.2/MODULE.bazel": "e2c890c8a515d6bca9c66d47718aa9e44b458fde64ec7204b8030bf2d349058c", -- "https://bcr.bazel.build/modules/bazel_lib/3.2.2/source.json": "9e84e115c20e14652c5c21401ae85ff4daa8702e265b5c0b3bf89353f17aa212", -+ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/MODULE.bazel": "cbc920fff00a55b2fd663d40850560a6fb17fd5ef3a9f3d6cbf0d32908dd8a40", -+ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/source.json": "7a9c2b06a113c4962d329bb68befdb90d0447788a19aabaa32b476a0458e7d28", - "https://bcr.bazel.build/modules/bazel_skylib/1.0.3/MODULE.bazel": "bcb0fd896384802d1ad283b4e4eb4d718eebd8cb820b0a2c3a347fb971afd9d8", - "https://bcr.bazel.build/modules/bazel_skylib/1.1.1/MODULE.bazel": "1add3e7d93ff2e6998f9e118022c84d163917d912f5afafb3058e3d2f1545b5e", - "https://bcr.bazel.build/modules/bazel_skylib/1.2.0/MODULE.bazel": "44fe84260e454ed94ad326352a698422dbe372b21a1ac9f3eab76eb531223686", -@@ -313,6 +314,8 @@ - "https://bcr.bazel.build/modules/windows_support/0.4.1/source.json": "bdb72e6717b00135e0a4a8512b3f6b626df0baea0bde74bf7312c7484f791287", - "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/MODULE.bazel": "b573395fe63aef4299ba095173e2f62ccfee5ad9bbf7acaa95dba73af9fc2b38", - "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/source.json": "3f3fbaeafecaf629877ad152a2c9def21f8d330d91aa94c5dc75bbb98c10b8b8", -+ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/MODULE.bazel": "ff7cf8a5334a96b7b9ffcbe73418d45929ce125efebac39c1ee7d14bd2ae7a9b", -+ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/source.json": "abf61f3c1630eab7b81722f710060b6f31e7c160511bcbfaad74bb4385883051", - "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/MODULE.bazel": "86130924115bf8dfc4aba286fb49a2ec8d32430a21771db9ccf5381acf0d10cc", - "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/source.json": "7aa4112eb2fbb99757e97cf0b7a4e2bdb3464b2b3cfb351bc807ffcb34bd349f", - "https://bcr.bazel.build/modules/zlib-ng/2.3.3/MODULE.bazel": "0afa781a5f354b4fd811a5c9a086e111daf50e09d8e98ab1cf7c46eeacc33aac", ++define void @store_low1_nxv8f16_splat(ptr %0) { ++; CHECK-LABEL: store_low1_nxv8f16_splat: ++; CHECK: // %bb.0: ++; CHECK-NEXT: fmov h0, #1.00000000 ++; CHECK-NEXT: str h0, [x0] ++; CHECK-NEXT: ret ++ %2 = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 1) ++ tail call void @llvm.masked.store.nxv8f16.p0( splat (half 0xH3C00), ptr align 2 %0, %2) ++ ret void ++} diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl index 88052c28512ca5..ea77ef6bd74546 100644 --- a/third_party/xla/third_party/llvm/workspace.bzl +++ b/third_party/xla/third_party/llvm/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") def repo(name): """Imports LLVM.""" - LLVM_COMMIT = "e56c2cefc3e7978de4a2d799fa3abec77d75af21" - LLVM_SHA256 = "affcad31c9455e95faca9d4de38e4a49b73c23bbf48b405d264ff10a6bfd3baf" + LLVM_COMMIT = "6ec87e7ec4b3ca6f665a02e05d0e5fe9c37cd6cc" + LLVM_SHA256 = "75300dbf9232d135c1ee2819216b7b7597fd9d36a90312c2b0c45938097d806f" tf_http_archive( name = name, diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch index 0d8bf613295b57..f589d2076e8bb6 100644 --- a/third_party/xla/third_party/shardy/temporary.patch +++ b/third_party/xla/third_party/shardy/temporary.patch @@ -1,1215 +1,734 @@ -diff --git a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir -index 5697a287..6ef5c695 100644 ---- a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir -+++ b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir -@@ -1,4 +1,4 @@ --// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=True' 2>&1 | FileCheck %s -+// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=true' 2>&1 | FileCheck %s - - !mesh_1_tensor_ui32 = !mpmd.mesh_tensor<"m1", tensor> - !mesh_1_tensor_1_ui32 = !mpmd.mesh_tensor<"m1", tensor<1xui32>> -diff --git a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir -index d914bd27..54d0ca3a 100644 ---- a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir -+++ b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir -@@ -1,4 +1,4 @@ --// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=True' 2>&1 | FileCheck %s -+// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=true' 2>&1 | FileCheck %s - - // Test that stablehlo.send and stablehlo.create_token ops (which produce - // !stablehlo.token types) are correctly partitioned into fragments by the diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch -index a9462de1..634ed2ef 100644 +index 634ed2ef..9da29843 100644 --- a/third_party/llvm/generated.patch +++ b/third_party/llvm/generated.patch -@@ -1,189 +1,664 @@ +@@ -1,664 +1,59 @@ Auto generated patch. Do not edit or delete it, even if empty. --diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp ----- a/clang/lib/CodeGen/CGCall.cpp --+++ b/clang/lib/CodeGen/CGCall.cpp --@@ -2804,17 +2804,10 @@ -- AddAttributesFromFunctionProtoType( -- getContext(), FuncAttrs, Fn->getType()->getAs()); -- if (AttrOnCallSite && Fn->isReplaceableGlobalAllocationFunction()) { --- // A sane operator new returns a non-aliasing pointer and does not --- // read or write accessible memory. --+ // A sane operator new returns a non-aliasing pointer. -- if (getCodeGenOpts().AssumeSaneOperatorNew && --- Fn->getDeclName().isAnyOperatorNew()) { --+ Fn->getDeclName().isAnyOperatorNew()) -- RetAttrs.addAttribute(llvm::Attribute::NoAlias); --- // FIXME: inaccessiblemem could cause issues if LTO makes the --- // previously inaccessible memory accessible after linking. --- FuncAttrs.addMemoryAttr( --- llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( --- llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); --- } -- } -- const CXXMethodDecl *MD = dyn_cast(Fn); -- const bool IsVirtualCall = MD && MD->isVirtual(); --diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGExprCXX.cpp b/clang/lib/CodeGen/CGExprCXX.cpp ----- a/clang/lib/CodeGen/CGExprCXX.cpp --+++ b/clang/lib/CodeGen/CGExprCXX.cpp --@@ -1391,9 +1391,19 @@ -- /// -- /// We model such elidable calls with the 'builtin' attribute. -- llvm::Function *Fn = dyn_cast(CalleePtr); --- if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn && --- Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) { --- CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); --+ if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn) { --+ if (Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) --+ CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); -+diff -ruN --strip-trailing-cr a/clang/include/clang/AST/TypeBase.h b/clang/include/clang/AST/TypeBase.h -+--- a/clang/include/clang/AST/TypeBase.h -++++ b/clang/include/clang/AST/TypeBase.h -+@@ -2270,7 +2270,7 @@ -+ /// increments towards the beginning. -+ /// Positive non-zero number represents the index + 1. -+ /// Zero means this is not substituted from an expansion. -+- unsigned PackIndex : 15; -++ unsigned PackIndex : 16; -+ }; -+ -+ class SubstPackTypeBitfields { -+diff -ruN --strip-trailing-cr a/clang/test/SemaCXX/make_integer_seq.cpp b/clang/test/SemaCXX/make_integer_seq.cpp -+--- a/clang/test/SemaCXX/make_integer_seq.cpp -++++ b/clang/test/SemaCXX/make_integer_seq.cpp -+@@ -50,3 +50,8 @@ -+ -+ __make_integer_seq<__make_integer_seq, int, 10> PR28494; // expected-note{{different template parameters}} -+ // expected-error@make_integer_seq.cpp:* {{template argument for template template parameter must be a class template or type alias template}} - + --+ // A sane operator new does not read or write accessible memory. --+ if (CGF.CGM.getCodeGenOpts().AssumeSaneOperatorNew && --+ CalleeDecl->getDeclName().isAnyOperatorNew()) { --+ // FIXME: inaccessiblemem could cause issues if LTO makes the --+ // previously inaccessible memory accessible after linking. --+ CallOrInvoke->setMemoryEffects( --+ llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( --+ llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); --+ } -++// The largest type pack whose pack indices fit in SubstTemplateTypeParmType. -++template using Expand = void(Ts...); -++template struct ExpandSeq { using type = Expand; }; -++using LargePack = __make_integer_seq::type; -+diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp -+--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp -++++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp -+@@ -13901,47 +13901,12 @@ -+ return SDValue(); -+ } -+ -+-/// Check if Mask defines a known constant set of enabled lanes, where only the -+-/// first N lanes are enabled. N is returned if so. -+-static uint64_t calculateConstantLowMaskLanes(SDValue Mask) { -+- // We expect masks for masked load/store to be i1 predicates. -+- if (Mask.getValueType().getScalarSizeInBits() != 1) -+- return 0; -+- -+- if (Mask.getOpcode() == ISD::BUILD_VECTOR) { -+- unsigned NumOnes = 0; -+- auto Op = Mask->op_begin(); -+- for (; Op != Mask->op_end(); ++Op) { -+- if (!isOneConstant(*Op)) -+- break; -+- ++NumOnes; -+- } -+- -+- if (NumOnes == 0) -+- return 0; -+- -+- for (; Op != Mask->op_end(); ++Op) -+- if (!isNullConstant(*Op)) -+- return 0; -+- return NumOnes; -+- } -+- -+- if (Mask.getOpcode() == ISD::GET_ACTIVE_LANE_MASK && -+- isNullConstant(Mask.getOperand(0)) && -+- isa(Mask.getOperand(1)) && -+- Mask.getOperand(1).getValueType().getSizeInBits() <= 64) -+- return Mask.getConstantOperandVal(1); -+- -+- return 0; -+-} -+- -+ SDValue DAGCombiner::visitMSTORE(SDNode *N) { -+ MaskedStoreSDNode *MST = cast(N); -+ SDValue Mask = MST->getMask(); -+ SDValue Chain = MST->getChain(); -+ SDValue Value = MST->getValue(); -+ SDValue Ptr = MST->getBasePtr(); -+- EVT VT = Value.getValueType(); +-diff -ruN --strip-trailing-cr a/clang/include/clang/AST/TypeBase.h b/clang/include/clang/AST/TypeBase.h +---- a/clang/include/clang/AST/TypeBase.h +-+++ b/clang/include/clang/AST/TypeBase.h +-@@ -2270,7 +2270,7 @@ +- /// increments towards the beginning. +- /// Positive non-zero number represents the index + 1. +- /// Zero means this is not substituted from an expansion. +-- unsigned PackIndex : 15; +-+ unsigned PackIndex : 16; +- }; +- +- class SubstPackTypeBitfields { +-diff -ruN --strip-trailing-cr a/clang/test/SemaCXX/make_integer_seq.cpp b/clang/test/SemaCXX/make_integer_seq.cpp +---- a/clang/test/SemaCXX/make_integer_seq.cpp +-+++ b/clang/test/SemaCXX/make_integer_seq.cpp +-@@ -50,3 +50,8 @@ +- +- __make_integer_seq<__make_integer_seq, int, 10> PR28494; // expected-note{{different template parameters}} +- // expected-error@make_integer_seq.cpp:* {{template argument for template template parameter must be a class template or type alias template}} +-+ +-+// The largest type pack whose pack indices fit in SubstTemplateTypeParmType. +-+template using Expand = void(Ts...); +-+template struct ExpandSeq { using type = Expand; }; +-+using LargePack = __make_integer_seq::type; +-diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +---- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +-+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +-@@ -13901,47 +13901,12 @@ +- return SDValue(); ++diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h ++--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h +++++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h ++@@ -837,7 +837,7 @@ + -+ // Zap masked stores with a zero mask. -+ if (ISD::isConstantSplatVectorAllZeros(Mask.getNode())) -+@@ -13964,46 +13929,21 @@ -+ } - } ++ SDValue ScalarizeVecRes_ADDRSPACECAST(SDNode *N); ++ SDValue ScalarizeVecRes_BITCAST(SDNode *N); ++- SDValue ScalarizeVecRes_BUILD_VECTOR(SDNode *N); +++ SDValue ScalarizeVecRes_BUILD_VECTOR_OR_SPLAT(SDNode *N); ++ SDValue ScalarizeVecRes_EXTRACT_SUBVECTOR(SDNode *N); ++ SDValue ScalarizeVecRes_FP_ROUND(SDNode *N); ++ SDValue ScalarizeVecRes_CONVERT_FROM_ARBITRARY_FP(SDNode *N); ++diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp ++--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp +++++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp ++@@ -63,7 +63,10 @@ ++ break; ++ case ISD::MERGE_VALUES: R = ScalarizeVecRes_MERGE_VALUES(N, ResNo);break; ++ case ISD::BITCAST: R = ScalarizeVecRes_BITCAST(N); break; ++- case ISD::BUILD_VECTOR: R = ScalarizeVecRes_BUILD_VECTOR(N); break; +++ case ISD::SPLAT_VECTOR: +++ case ISD::BUILD_VECTOR: +++ R = ScalarizeVecRes_BUILD_VECTOR_OR_SPLAT(N); +++ break; ++ case ISD::EXTRACT_SUBVECTOR: R = ScalarizeVecRes_EXTRACT_SUBVECTOR(N); break; ++ case ISD::FP_ROUND: R = ScalarizeVecRes_FP_ROUND(N); break; ++ case ISD::CONVERT_FROM_ARBITRARY_FP: ++@@ -475,10 +478,10 @@ ++ NewVT, Op); + } -- return RV; --diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/new_hot_cold.cpp b/clang/test/CodeGenCXX/new_hot_cold.cpp ----- a/clang/test/CodeGenCXX/new_hot_cold.cpp --+++ b/clang/test/CodeGenCXX/new_hot_cold.cpp --@@ -124,7 +124,7 @@ -+- // If this is a masked store with an all ones mask, we can use a unmasked -+- // store. -++ // If this is a masked load with an all ones mask, we can use a unmasked load. -+ // FIXME: Can we do this for indexed, compressing, or truncating stores? -+- if (MST->isUnindexed() && !MST->isCompressingStore() && -+- !MST->isTruncatingStore()) { -+- if (ISD::isConstantSplatVectorAllOnes(Mask.getNode())) -+- return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), -+- MST->getBasePtr(), MST->getPointerInfo(), -+- MST->getBaseAlign(), MST->getMemOperand()->getFlags(), -+- MST->getAAInfo()); -+- -+- // Convert a masked_store with constant getactivelanemask input mask to a -+- // standard store. -+- if (uint64_t Lanes = calculateConstantLowMaskLanes(Mask)) { -+- if (Lanes < VT.getVectorMinNumElements() && isPowerOf2_32(Lanes)) { -+- EVT SubVT = EVT::getVectorVT(*DAG.getContext(), -+- VT.getVectorElementType(), Lanes); -+- unsigned IsFast = 0; -+- if ((!LegalTypes || isTypeLegal(SubVT)) && -+- TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), -+- SubVT, MST->getAddressSpace(), -+- MST->getBaseAlign(), -+- MST->getMemOperand()->getFlags(), &IsFast) && -+- IsFast) { -+- SDLoc DL(N); -+- SDValue Ext = DAG.getExtractSubvector(DL, SubVT, MST->getValue(), 0); -+- return DAG.getStore(MST->getChain(), DL, Ext, MST->getBasePtr(), -+- MST->getPointerInfo(), MST->getBaseAlign(), -+- MST->getMemOperand()->getFlags(), -+- MST->getAAInfo()); -+- } -+- } -+- } -+- } -++ if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MST->isUnindexed() && -++ !MST->isCompressingStore() && !MST->isTruncatingStore()) -++ return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), -++ MST->getBasePtr(), MST->getPointerInfo(), -++ MST->getBaseAlign(), MST->getMemOperand()->getFlags(), -++ MST->getAAInfo()); - -- // CHECK-DAG: attributes [[ATTR_NOBUILTIN]] = { nobuiltin allocsize(0) {{.*}} } -- // CHECK-DAG: attributes [[ATTR_NOBUILTIN_NOTHROW]] = { nobuiltin nounwind allocsize(0) {{.*}} } ---// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } --+// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) } -- // CHECK-DAG: attributes [[ATTR_BUILTIN_CALL]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } ---// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } --+// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) } -- // CHECK-DAG: attributes [[ATTR_BUILTIN_NOTHROW_CALL]] = { builtin nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } --diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/operator-new.cpp b/clang/test/CodeGenCXX/operator-new.cpp ----- a/clang/test/CodeGenCXX/operator-new.cpp --+++ b/clang/test/CodeGenCXX/operator-new.cpp --@@ -27,5 +27,14 @@ -- } -+ // Try transforming N to an indexed store. -+ if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N)) -+ return SDValue(N, 0); -+ -+- if (MST->isTruncatingStore() && MST->isUnindexed() && VT.isInteger() && -++ if (MST->isTruncatingStore() && MST->isUnindexed() && -++ Value.getValueType().isInteger() && -+ (!isa(Value) || -+ !cast(Value)->isOpaque())) { -+ APInt TruncDemandedBits = -+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll -+--- a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll -++++ b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll -+@@ -139,7 +139,14 @@ -+ define void @store_low2_v4f64(ptr %p, <4 x double> %a) vscale_range(2, 2) { -+ ; CHECK-LABEL: store_low2_v4f64: -+ ; CHECK: // %bb.0: -+-; CHECK-NEXT: str q0, [x0] -++; CHECK-NEXT: // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1 -++; CHECK-NEXT: ptrue p0.d, vl2 -++; CHECK-NEXT: ptrue p1.h, vl2 -++; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1 -++; CHECK-NEXT: splice z0.d, p0, { z0.d, z1.d } -++; CHECK-NEXT: punpklo p0.h, p1.b -++; CHECK-NEXT: punpklo p0.h, p0.b -++; CHECK-NEXT: st1d { z0.d }, p0, [x0] -+ ; CHECK-NEXT: ret -+ %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) -+ tail call void @llvm.masked.store(<4 x double> %a, ptr align 1 %p, <4 x i1> %m) -+@@ -149,7 +156,13 @@ -+ define void @store_low2_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { -+ ; CHECK-LABEL: store_low2_v4i32: -+ ; CHECK: // %bb.0: -+-; CHECK-NEXT: str d0, [x0] -++; CHECK-NEXT: ptrue p0.h, vl2 -++; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 -++; CHECK-NEXT: mov z1.h, p0/z, #-1 // =0xffffffffffffffff -++; CHECK-NEXT: ptrue p0.s -++; CHECK-NEXT: sshll v1.4s, v1.4h, #0 -++; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 -++; CHECK-NEXT: st1w { z0.s }, p1, [x0] -+ ; CHECK-NEXT: ret -+ %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) -+ tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> %m) -+@@ -159,7 +172,12 @@ -+ define void @store_low2c_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { -+ ; CHECK-LABEL: store_low2c_v4i32: -+ ; CHECK: // %bb.0: -+-; CHECK-NEXT: str d0, [x0] -++; CHECK-NEXT: adrp x8, .LCPI12_0 -++; CHECK-NEXT: ptrue p0.s -++; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 -++; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI12_0] -++; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 -++; CHECK-NEXT: st1w { z0.s }, p1, [x0] -+ ; CHECK-NEXT: ret -+ tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> ) +--/// Check if Mask defines a known constant set of enabled lanes, where only the +--/// first N lanes are enabled. N is returned if so. +--static uint64_t calculateConstantLowMaskLanes(SDValue Mask) { +-- // We expect masks for masked load/store to be i1 predicates. +-- if (Mask.getValueType().getScalarSizeInBits() != 1) +-- return 0; +-- +-- if (Mask.getOpcode() == ISD::BUILD_VECTOR) { +-- unsigned NumOnes = 0; +-- auto Op = Mask->op_begin(); +-- for (; Op != Mask->op_end(); ++Op) { +-- if (!isOneConstant(*Op)) +-- break; +-- ++NumOnes; +-- } +-- +-- if (NumOnes == 0) +-- return 0; +-- +-- for (; Op != Mask->op_end(); ++Op) +-- if (!isNullConstant(*Op)) +-- return 0; +-- return NumOnes; +-- } +-- +-- if (Mask.getOpcode() == ISD::GET_ACTIVE_LANE_MASK && +-- isNullConstant(Mask.getOperand(0)) && +-- isa(Mask.getOperand(1)) && +-- Mask.getOperand(1).getValueType().getSizeInBits() <= 64) +-- return Mask.getConstantOperandVal(1); +-- +-- return 0; +--} +-- +- SDValue DAGCombiner::visitMSTORE(SDNode *N) { +- MaskedStoreSDNode *MST = cast(N); +- SDValue Mask = MST->getMask(); +- SDValue Chain = MST->getChain(); +- SDValue Value = MST->getValue(); +- SDValue Ptr = MST->getBasePtr(); +-- EVT VT = Value.getValueType(); +- +- // Zap masked stores with a zero mask. +- if (ISD::isConstantSplatVectorAllZeros(Mask.getNode())) +-@@ -13964,46 +13929,21 @@ +- } +- } +- +-- // If this is a masked store with an all ones mask, we can use a unmasked +-- // store. +-+ // If this is a masked load with an all ones mask, we can use a unmasked load. +- // FIXME: Can we do this for indexed, compressing, or truncating stores? +-- if (MST->isUnindexed() && !MST->isCompressingStore() && +-- !MST->isTruncatingStore()) { +-- if (ISD::isConstantSplatVectorAllOnes(Mask.getNode())) +-- return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), +-- MST->getBasePtr(), MST->getPointerInfo(), +-- MST->getBaseAlign(), MST->getMemOperand()->getFlags(), +-- MST->getAAInfo()); +-- +-- // Convert a masked_store with constant getactivelanemask input mask to a +-- // standard store. +-- if (uint64_t Lanes = calculateConstantLowMaskLanes(Mask)) { +-- if (Lanes < VT.getVectorMinNumElements() && isPowerOf2_32(Lanes)) { +-- EVT SubVT = EVT::getVectorVT(*DAG.getContext(), +-- VT.getVectorElementType(), Lanes); +-- unsigned IsFast = 0; +-- if ((!LegalTypes || isTypeLegal(SubVT)) && +-- TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), +-- SubVT, MST->getAddressSpace(), +-- MST->getBaseAlign(), +-- MST->getMemOperand()->getFlags(), &IsFast) && +-- IsFast) { +-- SDLoc DL(N); +-- SDValue Ext = DAG.getExtractSubvector(DL, SubVT, MST->getValue(), 0); +-- return DAG.getStore(MST->getChain(), DL, Ext, MST->getBasePtr(), +-- MST->getPointerInfo(), MST->getBaseAlign(), +-- MST->getMemOperand()->getFlags(), +-- MST->getAAInfo()); +-- } +-- } +-- } +-- } +-+ if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MST->isUnindexed() && +-+ !MST->isCompressingStore() && !MST->isTruncatingStore()) +-+ return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), +-+ MST->getBasePtr(), MST->getPointerInfo(), +-+ MST->getBaseAlign(), MST->getMemOperand()->getFlags(), +-+ MST->getAAInfo()); +- +- // Try transforming N to an indexed store. +- if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N)) +- return SDValue(N, 0); +- +-- if (MST->isTruncatingStore() && MST->isUnindexed() && VT.isInteger() && +-+ if (MST->isTruncatingStore() && MST->isUnindexed() && +-+ Value.getValueType().isInteger() && +- (!isa(Value) || +- !cast(Value)->isOpaque())) { +- APInt TruncDemandedBits = ++-SDValue DAGTypeLegalizer::ScalarizeVecRes_BUILD_VECTOR(SDNode *N) { +++SDValue DAGTypeLegalizer::ScalarizeVecRes_BUILD_VECTOR_OR_SPLAT(SDNode *N) { ++ EVT EltVT = N->getValueType(0).getVectorElementType(); ++ SDValue InOp = N->getOperand(0); ++- // The BUILD_VECTOR operands may be of wider element types and +++ // The BUILD_VECTOR / SPLAT operands may be of wider element types and ++ // we may need to truncate them back to the requested return type. ++ if (EltVT.isInteger()) ++ return DAG.getNode(ISD::TRUNCATE, SDLoc(N), EltVT, InOp); + diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll + --- a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll + +++ b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll +-@@ -139,7 +139,14 @@ +- define void @store_low2_v4f64(ptr %p, <4 x double> %a) vscale_range(2, 2) { +- ; CHECK-LABEL: store_low2_v4f64: +- ; CHECK: // %bb.0: +--; CHECK-NEXT: str q0, [x0] +-+; CHECK-NEXT: // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1 +-+; CHECK-NEXT: ptrue p0.d, vl2 +-+; CHECK-NEXT: ptrue p1.h, vl2 +-+; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1 +-+; CHECK-NEXT: splice z0.d, p0, { z0.d, z1.d } +-+; CHECK-NEXT: punpklo p0.h, p1.b +-+; CHECK-NEXT: punpklo p0.h, p0.b +-+; CHECK-NEXT: st1d { z0.d }, p0, [x0] +- ; CHECK-NEXT: ret +- %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) +- tail call void @llvm.masked.store(<4 x double> %a, ptr align 1 %p, <4 x i1> %m) +-@@ -149,7 +156,13 @@ +- define void @store_low2_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { +- ; CHECK-LABEL: store_low2_v4i32: +- ; CHECK: // %bb.0: +--; CHECK-NEXT: str d0, [x0] +-+; CHECK-NEXT: ptrue p0.h, vl2 +-+; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 +-+; CHECK-NEXT: mov z1.h, p0/z, #-1 // =0xffffffffffffffff +-+; CHECK-NEXT: ptrue p0.s +-+; CHECK-NEXT: sshll v1.4s, v1.4h, #0 +-+; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 +-+; CHECK-NEXT: st1w { z0.s }, p1, [x0] +- ; CHECK-NEXT: ret +- %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) +- tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> %m) +-@@ -159,7 +172,12 @@ +- define void @store_low2c_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { +- ; CHECK-LABEL: store_low2c_v4i32: +- ; CHECK: // %bb.0: +--; CHECK-NEXT: str d0, [x0] +-+; CHECK-NEXT: adrp x8, .LCPI12_0 +-+; CHECK-NEXT: ptrue p0.s +-+; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 +-+; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI12_0] +-+; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 +-+; CHECK-NEXT: st1w { z0.s }, p1, [x0] +- ; CHECK-NEXT: ret +- tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> ) +- ret void +-@@ -168,7 +186,11 @@ +- define void @store_low2_nxv4f64(ptr %p, %a) { +- ; CHECK-LABEL: store_low2_nxv4f64: +- ; CHECK: // %bb.0: +--; CHECK-NEXT: str q0, [x0] +-+; CHECK-NEXT: ptrue p0.s, vl2 +-+; CHECK-NEXT: punpkhi p1.h, p0.b +-+; CHECK-NEXT: punpklo p0.h, p0.b +-+; CHECK-NEXT: st1d { z1.d }, p1, [x0, #1, mul vl] +-+; CHECK-NEXT: st1d { z0.d }, p0, [x0] +- ; CHECK-NEXT: ret +- %m = tail call @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 2) +- tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) +-@@ -203,10 +225,13 @@ +- define void @store_low2_v8i16(ptr %p, <8 x i16> %a) vscale_range(2, 2) { +- ; CHECK-LABEL: store_low2_v8i16: +- ; CHECK: // %bb.0: +-+; CHECK-NEXT: ptrue p0.b, vl2 +- ; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 +--; CHECK-NEXT: ptrue p0.s, vl2 +--; CHECK-NEXT: uunpklo z0.s, z0.h +--; CHECK-NEXT: st1h { z0.s }, p0, [x0] +-+; CHECK-NEXT: mov z1.b, p0/z, #-1 // =0xffffffffffffffff +-+; CHECK-NEXT: ptrue p0.h, vl8 +-+; CHECK-NEXT: sshll v1.8h, v1.8b, #0 +-+; CHECK-NEXT: cmpne p1.h, p0/z, z1.h, #0 +-+; CHECK-NEXT: st1h { z0.h }, p1, [x0] +- ; CHECK-NEXT: ret +- %m = tail call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 0, i32 2) +- tail call void @llvm.masked.store(<8 x i16> %a, ptr align 1 %p, <8 x i1> %m) +-@@ -216,9 +241,8 @@ +- define void @store_low2_nxv8i16(ptr %p, %a) { +- ; CHECK-LABEL: store_low2_nxv8i16: +- ; CHECK: // %bb.0: +--; CHECK-NEXT: mov z1.h, z0.h[1] +--; CHECK-NEXT: str h0, [x0] +--; CHECK-NEXT: str h1, [x0, #2] +-+; CHECK-NEXT: ptrue p0.h, vl2 +-+; CHECK-NEXT: st1h { z0.h }, p0, [x0] +- ; CHECK-NEXT: ret +- %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 2) +- tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) +-@@ -239,7 +263,8 @@ +- define void @store_low4_nxv8i16(ptr %p, %a) { +- ; CHECK-LABEL: store_low4_nxv8i16: +- ; CHECK: // %bb.0: +--; CHECK-NEXT: str d0, [x0] +-+; CHECK-NEXT: ptrue p0.h, vl4 +-+; CHECK-NEXT: st1h { z0.h }, p0, [x0] +- ; CHECK-NEXT: ret +- %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 4) ++@@ -245,3 +245,14 @@ + tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) +-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll +---- a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll +-+++ b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll +-@@ -6,43 +6,40 @@ +- ; CHECK-LABEL: f0: +- ; CHECK: // %bb.0: // %b0 +- ; CHECK-NEXT: { +--; CHECK-NEXT: r0 = add(r1,r0) +-+; CHECK-NEXT: [[RADD:[r0-9]+]] = add([[R1:[r0-9]+]],[[R0:[r0-9]+]]) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: allocframe(r29,#128):raw +-+; CHECK-NEXT: [[RLCPI:[r0-9]+]] = ##.LCPI0_0 +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: r2 = add(r0,#136) +-+; CHECK-NEXT: [[RNEG:[r0-9]+]] = #-1 +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: r29 = and(r29,#-128) +-+; CHECK-NEXT: [[RADD136:[r0-9]+]] = add([[RADD]],#136) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: v1 = vmem(r0+#1) +-+; CHECK-NEXT: [[V1:[v0-9]+]] = vmem([[RLCPI]]+#0) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: v0 = vmemu(r2+#0) +-+; CHECK-NEXT: [[RADD128:[r0-9]+]] = add([[RADD]],#128) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: r4 = add(r29,#0) +-+; CHECK-NEXT: [[Q0:[q0-9]+]] = vand([[V1]],[[RNEG]]) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: v0.w = vadd(v1.w,v0.w) +-+; CHECK-NEXT: [[V2:[v0-9]+]] = vmemu([[RADD136]]+#0) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: vmem(r4+#0) = v0 +-+; CHECK-NEXT: [[V3:[v0-9]+]] = vmem([[RADD]]+#1) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: r2 = memw(r4+#0) +-+; CHECK-NEXT: [[V0:[v0-9]+]].w = vadd([[V3]].w,[[V2]].w) +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: r3 = memw(r4+#4) +-+; CHECK-NEXT: if ([[Q0]]) vmem([[RADD128]]+#0) = [[V0]] +- ; CHECK-NEXT: } +- ; CHECK-NEXT: { +--; CHECK-NEXT: memd(r0+#128) = r3:2 +--; CHECK-NEXT: } +--; CHECK-NEXT: { +--; CHECK-NEXT: r31:30 = dealloc_return(r30):raw +-+; CHECK-NEXT: jumpr r31 +- ; CHECK-NEXT: } +- b0: +- %v0 = add i32 %a1, 128 +-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll +---- a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll +-+++ b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll +-@@ -29,12 +29,20 @@ +- define void @kmovkr_2() { +- ; AVX512-LABEL: kmovkr_2: +- ; AVX512: # %bb.0: # %alloca_21 +--; AVX512-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] +-+; AVX512-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] +-+; AVX512-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] +-+; AVX512-NEXT: kmovw %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xf8,0x92,0xc8] +-+; AVX512-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] +-+; AVX512-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] +- ; AVX512-NEXT: retq # encoding: [0xc3] +- ; +- ; AVX512BW-LABEL: kmovkr_2: +- ; AVX512BW: # %bb.0: # %alloca_21 +--; AVX512BW-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] +-+; AVX512BW-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] +-+; AVX512BW-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] +-+; AVX512BW-NEXT: kmovd %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xfb,0x92,0xc8] +-+; AVX512BW-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] +-+; AVX512BW-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] +- ; AVX512BW-NEXT: retq # encoding: [0xc3] +- alloca_21: +- call void @llvm.masked.store.v4f32.p0(<4 x float> zeroinitializer, ptr null, i32 1, <4 x i1> ) +-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll +---- a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll +-+++ b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll +-@@ -1,21 +0,0 @@ +--; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py +--; RUN: llc < %s -mtriple=x86_64-linux-gnu -mattr=+avx,+avx2,+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl | FileCheck %s +-- +--define fastcc void @render_event(<12 x double> %0) { +--; CHECK-LABEL: render_event: +--; CHECK: # %bb.0: +--; CHECK-NEXT: vmovlhps {{.*#+}} xmm6 = xmm6[0],xmm7[0] +--; CHECK-NEXT: vmovlhps {{.*#+}} xmm4 = xmm4[0],xmm5[0] +--; CHECK-NEXT: vinsertf128 $1, %xmm6, %ymm4, %ymm4 +--; CHECK-NEXT: vmovlhps {{.*#+}} xmm2 = xmm2[0],xmm3[0] +--; CHECK-NEXT: vmovlhps {{.*#+}} xmm0 = xmm0[0],xmm1[0] +--; CHECK-NEXT: vinsertf128 $1, %xmm2, %ymm0, %ymm0 +--; CHECK-NEXT: vinsertf64x4 $1, %ymm4, %zmm0, %zmm0 +--; CHECK-NEXT: vmovsd {{.*#+}} xmm1 = mem[0],zero +--; CHECK-NEXT: vmovsd %xmm1, 64 +--; CHECK-NEXT: vmovups %zmm0, 0 +--; CHECK-NEXT: vzeroupper +--; CHECK-NEXT: retq +-- call void @llvm.masked.store.v12f64.p0(<12 x double> %0, ptr null, <12 x i1> ) +-- ret void +--} +-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr179489.ll b/llvm/test/CodeGen/X86/pr179489.ll +---- a/llvm/test/CodeGen/X86/pr179489.ll +-+++ b/llvm/test/CodeGen/X86/pr179489.ll +-@@ -9,12 +9,14 @@ +- ; X86-LABEL: foo: +- ; X86: # %bb.0: # %entry +- ; X86-NEXT: movl {{[0-9]+}}(%esp), %eax +--; X86-NEXT: vpextrb $1, %xmm0, (%eax) +-+; X86-NEXT: vmovw %xmm0, %ecx +-+; X86-NEXT: movb %ch, (%eax) +- ; X86-NEXT: retl +- ; +- ; X64-LABEL: foo: +- ; X64: # %bb.0: # %entry +--; X64-NEXT: vpextrb $1, %xmm0, (%rdi) +-+; X64-NEXT: vmovw %xmm0, %eax +-+; X64-NEXT: movb %ah, (%rdi) +- ; X64-NEXT: retq +- entry: +- %i3 = shufflevector <8 x i16> %arg, <8 x i16> , <8 x i32> +-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr46532.ll b/llvm/test/CodeGen/X86/pr46532.ll +---- a/llvm/test/CodeGen/X86/pr46532.ll +-+++ b/llvm/test/CodeGen/X86/pr46532.ll +-@@ -7,6 +7,8 @@ +- ; CHECK-NEXT: movq (%rax), %rax +- ; CHECK-NEXT: vxorps %xmm0, %xmm0, %xmm0 +- ; CHECK-NEXT: vmovaps %xmm0, 32(%rax) +-+; CHECK-NEXT: vmovsd {{.*#+}} xmm0 = [4294967295,4294967295,0,0] +-+; CHECK-NEXT: vmaskmovps %ymm0, %ymm0, (%rax) +- while.1.body.preheader: +- %0 = load ptr, ptr undef, align 8, !invariant.load !0, !dereferenceable !1, !align !2 +- %1 = getelementptr inbounds i8, ptr %0, i64 32 +-diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td +---- a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td +-+++ b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td +-@@ -1,71 +0,0 @@ +--// RUN: split-file %s %t +--// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ +--// RUN: | FileCheck %t/valid.td +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-bug.td 2>&1 \ +--// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=user-sgpr-init16-bug --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-priority.td 2>&1 \ +--// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=required-export-priority --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-timing.td 2>&1 \ +--// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=gfx125x-lowest-rate-wmma --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-banks.td 2>&1 \ +--// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=ldsbankcount32 --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-lds-size.td 2>&1 \ +--// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=addressablelocalmemorysize65536 --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/capability.td 2>&1 \ +--// RUN: | FileCheck %t/capability.td -DFILE=%t/capability.td --implicit-check-not="error:" +-- +--//--- common.td +--include "AMDGPU.td" +-- +--def EmptyGPU : AMDGPUProcessorModel<"gfx999", NoSchedModel, [], [9, 9, 9]>; +-- +--//--- valid.td +--include "common.td" +--// Check the real generic targets, including gfx11's combined workarounds, +--// gfx12-5's fixup and WMMA timing, and gfx9-4's LDS size and bank count. +--// CHECK: GK_GFX9_4_GENERIC +--// CHECK: GK_GFX11_GENERIC +--// CHECK: GK_GFX12_5_GENERIC +-- +--//--- missing-bug.td +--include "common.td" +--defvar CheckedFeature = FeatureUserSGPRInit16Bug; +--include "missing-generic.td" +-- +--//--- missing-priority.td +--include "common.td" +--defvar CheckedFeature = FeatureRequiredExportPriority; +--include "missing-generic.td" +-- +--//--- missing-timing.td +--include "common.td" +--defvar CheckedFeature = FeatureGFX125xLowestRateWMMA; +--include "missing-generic.td" +-- +--//--- missing-banks.td +--include "common.td" +--defvar CheckedFeature = FeatureLDSBankCount32; +--include "missing-generic.td" +-- +--//--- missing-lds-size.td +--include "common.td" +--defvar CheckedFeature = FeatureAddressableLocalMemorySize65536; +--include "missing-generic.td" +-- +--//--- missing-generic.td +--// Classified backend-only features still need support from a covered GPU. +--// Included files use the native path separator after the include directory. +--// CHECK: [[DIR]]{{[/\\]}}missing-generic.td:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature '[[FEATURE]]' not supported by any covered GPU +--def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [CheckedFeature], [9, 9, 0]> { +-- let CoveredGPUs = [EmptyGPU]; +--} +-- +--//--- capability.td +--include "common.td" +--def WMMAGPU : AMDGPUProcessorModel<"gfx998", NoSchedModel, +-- [FeatureWMMA256bInsts], [9, 9, 8]>; +--// Instruction support keeps the default all-covered-GPUs rule. +--// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature 'wmma-256b-insts' not supported by covered GPU 'gfx999' +--def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [FeatureWMMA256bInsts], [9, 9, 0]> { +-- let CoveredGPUs = [WMMAGPU, EmptyGPU]; +--} +-diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td +---- a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td +-+++ b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td +-@@ -1,33 +1,66 @@ +- // RUN: split-file %s %t +--// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ +-+// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/valid.td \ +- // RUN: | FileCheck %t/valid.td +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ +- // RUN: | FileCheck %t/missing-any.td -DFILE=%t/missing-any.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ +- // RUN: | FileCheck %t/missing-all.td -DFILE=%t/missing-all.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ +- // RUN: | FileCheck %t/implied-all.td -DFILE=%t/implied-all.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ +- // RUN: | FileCheck %t/missing-private.td -DFILE=%t/missing-private.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-private.td 2>&1 \ +-+// RUN: | FileCheck %t/implied-private.td -DFILE=%t/implied-private.td --implicit-check-not="error:" +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ +- // RUN: | FileCheck %t/numeric-mismatch.td -DFILE=%t/numeric-mismatch.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ +- // RUN: | FileCheck %t/empty-members.td -DFILE=%t/empty-members.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ +- // RUN: | FileCheck %t/self-member.td -DFILE=%t/self-member.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ +- // RUN: | FileCheck %t/generic-member.td -DFILE=%t/generic-member.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ +- // RUN: | FileCheck %t/pseudo-member.td -DFILE=%t/pseudo-member.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ +- // RUN: | FileCheck %t/unregistered-member.td -DFILE=%t/unregistered-member.td --implicit-check-not="error:" +--// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ +-+// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ +- // RUN: | FileCheck %t/r600-member.td -DFILE=%t/r600-member.td --implicit-check-not="error:" +- +- //--- common.td +--include "llvm/Target/Target.td" +-+// Exercise the validation rules with small synthetic targets. The real targets +-+// are validated when AMDGPUTargetParserDef.inc is generated during the build. +-+// Stand-ins for the Target.td classes used by AMDGPUTargetParser.td, limited to +-+// the fields the backend reads. +-+class InlineBehavior; +-+def InlineIgnore : InlineBehavior; +-+def InlineDefault : InlineBehavior; +-+class SubtargetFeature i = [], +-+ InlineBehavior ib = InlineDefault> { +-+ string Name = n; +-+ string FieldName = f; +-+ string Value = v; +-+ string Desc = d; +-+ list Implies = i; +-+ InlineBehavior InlineBehavior = ib; +-+} +-+ +-+class SchedMachineModel; +-+def NoSchedModel : SchedMachineModel; +-+class ProcessorItineraries; +-+def NoItineraries : ProcessorItineraries; +-+ +-+class Processor f> { +-+ string Name = n; +-+ ProcessorItineraries ProcItin = pi; +-+ list Features = f; +-+} +-+class ProcessorModel f> +-+ : Processor { +-+ SchedMachineModel SchedModel = m; +-+} +-+ +- include "AMDGPUTargetParser.td" +--def MyTarget : Target; +- +- def FeatureAll : SubtargetFeature<"all", "HasAll", "true", "">; +- def FeatureAnyA : AMDGPUGenericAnyFeature<"any-a", "HasAnyA", "true", "">; +-@@ -37,7 +70,13 @@ +- def FeatureAnySmall : NumericFeature<4>; +- def FeatureAnyB : NumericFeature<8>; +- def FeatureAnyLarge : NumericFeature<16>; +--def FeaturePrivate : AMDGPUGenericAnyFeature<"private", "HasPrivate", "true", "">; +-+// Classification inherited through a subclass also applies to backend-only +-+// features, which must participate in validation without being emitted. +-+class PrivateFeature : AMDGPUGenericAnyFeature< +-+ "private", "HasPrivate", "true", "">; +-+def FeaturePrivate : PrivateFeature; +-+def FeatureImpliesPrivate : SubtargetFeature< +-+ "implies-private", "HasImpliesPrivate", "true", "", [FeaturePrivate]>; +- def FeatureImpliesA : SubtargetFeature< +- "implies-a", "HasImpliesA", "true", "", [FeatureAnyA]>; +- def FeatureImpliesB : SubtargetFeature< +-@@ -55,7 +94,7 @@ +- def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, +- [FeatureAll, FeatureAnyA, FeatureAnySmall], [9, 0, 0]>; +- def GFX901 : AMDGPUProcessorModel<"gfx901", NoSchedModel, +-- [FeatureAll, FeatureImpliesB, FeaturePrivate], [9, 0, 1]>; +-+ [FeatureAll, FeatureImpliesB, FeatureImpliesPrivate], [9, 0, 1]>; +- // Different members supply different marked features. Follow Implies on both +- // the generic and the members, and check marked backend-only features too. +- def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, +-@@ -103,6 +142,15 @@ +- let CoveredGPUs = [GFX900]; + ret void -+@@ -168,7 +186,11 @@ -+ define void @store_low2_nxv4f64(ptr %p, %a) { -+ ; CHECK-LABEL: store_low2_nxv4f64: -+ ; CHECK: // %bb.0: -+-; CHECK-NEXT: str q0, [x0] -++; CHECK-NEXT: ptrue p0.s, vl2 -++; CHECK-NEXT: punpkhi p1.h, p0.b -++; CHECK-NEXT: punpklo p0.h, p0.b -++; CHECK-NEXT: st1d { z1.d }, p1, [x0, #1, mul vl] -++; CHECK-NEXT: st1d { z0.d }, p0, [x0] -+ ; CHECK-NEXT: ret -+ %m = tail call @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 2) -+ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) -+@@ -203,10 +225,13 @@ -+ define void @store_low2_v8i16(ptr %p, <8 x i16> %a) vscale_range(2, 2) { -+ ; CHECK-LABEL: store_low2_v8i16: -+ ; CHECK: // %bb.0: -++; CHECK-NEXT: ptrue p0.b, vl2 -+ ; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 -+-; CHECK-NEXT: ptrue p0.s, vl2 -+-; CHECK-NEXT: uunpklo z0.s, z0.h -+-; CHECK-NEXT: st1h { z0.s }, p0, [x0] -++; CHECK-NEXT: mov z1.b, p0/z, #-1 // =0xffffffffffffffff -++; CHECK-NEXT: ptrue p0.h, vl8 -++; CHECK-NEXT: sshll v1.8h, v1.8b, #0 -++; CHECK-NEXT: cmpne p1.h, p0/z, z1.h, #0 -++; CHECK-NEXT: st1h { z0.h }, p1, [x0] -+ ; CHECK-NEXT: ret -+ %m = tail call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 0, i32 2) -+ tail call void @llvm.masked.store(<8 x i16> %a, ptr align 1 %p, <8 x i1> %m) -+@@ -216,9 +241,8 @@ -+ define void @store_low2_nxv8i16(ptr %p, %a) { -+ ; CHECK-LABEL: store_low2_nxv8i16: -+ ; CHECK: // %bb.0: -+-; CHECK-NEXT: mov z1.h, z0.h[1] -+-; CHECK-NEXT: str h0, [x0] -+-; CHECK-NEXT: str h1, [x0, #2] -++; CHECK-NEXT: ptrue p0.h, vl2 -++; CHECK-NEXT: st1h { z0.h }, p0, [x0] -+ ; CHECK-NEXT: ret -+ %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 2) -+ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) -+@@ -239,7 +263,8 @@ -+ define void @store_low4_nxv8i16(ptr %p, %a) { -+ ; CHECK-LABEL: store_low4_nxv8i16: -+ ; CHECK: // %bb.0: -+-; CHECK-NEXT: str d0, [x0] -++; CHECK-NEXT: ptrue p0.h, vl4 -++; CHECK-NEXT: st1h { z0.h }, p0, [x0] -+ ; CHECK-NEXT: ret -+ %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 4) -+ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) -+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll -+--- a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll -++++ b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll -+@@ -6,43 +6,40 @@ -+ ; CHECK-LABEL: f0: -+ ; CHECK: // %bb.0: // %b0 -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: r0 = add(r1,r0) -++; CHECK-NEXT: [[RADD:[r0-9]+]] = add([[R1:[r0-9]+]],[[R0:[r0-9]+]]) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: allocframe(r29,#128):raw -++; CHECK-NEXT: [[RLCPI:[r0-9]+]] = ##.LCPI0_0 -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: r2 = add(r0,#136) -++; CHECK-NEXT: [[RNEG:[r0-9]+]] = #-1 -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: r29 = and(r29,#-128) -++; CHECK-NEXT: [[RADD136:[r0-9]+]] = add([[RADD]],#136) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: v1 = vmem(r0+#1) -++; CHECK-NEXT: [[V1:[v0-9]+]] = vmem([[RLCPI]]+#0) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: v0 = vmemu(r2+#0) -++; CHECK-NEXT: [[RADD128:[r0-9]+]] = add([[RADD]],#128) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: r4 = add(r29,#0) -++; CHECK-NEXT: [[Q0:[q0-9]+]] = vand([[V1]],[[RNEG]]) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: v0.w = vadd(v1.w,v0.w) -++; CHECK-NEXT: [[V2:[v0-9]+]] = vmemu([[RADD136]]+#0) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: vmem(r4+#0) = v0 -++; CHECK-NEXT: [[V3:[v0-9]+]] = vmem([[RADD]]+#1) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: r2 = memw(r4+#0) -++; CHECK-NEXT: [[V0:[v0-9]+]].w = vadd([[V3]].w,[[V2]].w) -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: r3 = memw(r4+#4) -++; CHECK-NEXT: if ([[Q0]]) vmem([[RADD128]]+#0) = [[V0]] -+ ; CHECK-NEXT: } -+ ; CHECK-NEXT: { -+-; CHECK-NEXT: memd(r0+#128) = r3:2 -+-; CHECK-NEXT: } -+-; CHECK-NEXT: { -+-; CHECK-NEXT: r31:30 = dealloc_return(r30):raw -++; CHECK-NEXT: jumpr r31 -+ ; CHECK-NEXT: } -+ b0: -+ %v0 = add i32 %a1, 128 -+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll -+--- a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll -++++ b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll -+@@ -29,12 +29,20 @@ -+ define void @kmovkr_2() { -+ ; AVX512-LABEL: kmovkr_2: -+ ; AVX512: # %bb.0: # %alloca_21 -+-; AVX512-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] -++; AVX512-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] -++; AVX512-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] -++; AVX512-NEXT: kmovw %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xf8,0x92,0xc8] -++; AVX512-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] -++; AVX512-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] -+ ; AVX512-NEXT: retq # encoding: [0xc3] -+ ; -+ ; AVX512BW-LABEL: kmovkr_2: -+ ; AVX512BW: # %bb.0: # %alloca_21 -+-; AVX512BW-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] -++; AVX512BW-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] -++; AVX512BW-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] -++; AVX512BW-NEXT: kmovd %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xfb,0x92,0xc8] -++; AVX512BW-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] -++; AVX512BW-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] -+ ; AVX512BW-NEXT: retq # encoding: [0xc3] -+ alloca_21: -+ call void @llvm.masked.store.v4f32.p0(<4 x float> zeroinitializer, ptr null, i32 1, <4 x i1> ) -+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll -+--- a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll -++++ b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll -+@@ -1,21 +0,0 @@ -+-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py -+-; RUN: llc < %s -mtriple=x86_64-linux-gnu -mattr=+avx,+avx2,+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl | FileCheck %s -+- -+-define fastcc void @render_event(<12 x double> %0) { -+-; CHECK-LABEL: render_event: -+-; CHECK: # %bb.0: -+-; CHECK-NEXT: vmovlhps {{.*#+}} xmm6 = xmm6[0],xmm7[0] -+-; CHECK-NEXT: vmovlhps {{.*#+}} xmm4 = xmm4[0],xmm5[0] -+-; CHECK-NEXT: vinsertf128 $1, %xmm6, %ymm4, %ymm4 -+-; CHECK-NEXT: vmovlhps {{.*#+}} xmm2 = xmm2[0],xmm3[0] -+-; CHECK-NEXT: vmovlhps {{.*#+}} xmm0 = xmm0[0],xmm1[0] -+-; CHECK-NEXT: vinsertf128 $1, %xmm2, %ymm0, %ymm0 -+-; CHECK-NEXT: vinsertf64x4 $1, %ymm4, %zmm0, %zmm0 -+-; CHECK-NEXT: vmovsd {{.*#+}} xmm1 = mem[0],zero -+-; CHECK-NEXT: vmovsd %xmm1, 64 -+-; CHECK-NEXT: vmovups %zmm0, 0 -+-; CHECK-NEXT: vzeroupper -+-; CHECK-NEXT: retq -+- call void @llvm.masked.store.v12f64.p0(<12 x double> %0, ptr null, <12 x i1> ) -+- ret void -+-} -+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr179489.ll b/llvm/test/CodeGen/X86/pr179489.ll -+--- a/llvm/test/CodeGen/X86/pr179489.ll -++++ b/llvm/test/CodeGen/X86/pr179489.ll -+@@ -9,12 +9,14 @@ -+ ; X86-LABEL: foo: -+ ; X86: # %bb.0: # %entry -+ ; X86-NEXT: movl {{[0-9]+}}(%esp), %eax -+-; X86-NEXT: vpextrb $1, %xmm0, (%eax) -++; X86-NEXT: vmovw %xmm0, %ecx -++; X86-NEXT: movb %ch, (%eax) -+ ; X86-NEXT: retl -+ ; -+ ; X64-LABEL: foo: -+ ; X64: # %bb.0: # %entry -+-; X64-NEXT: vpextrb $1, %xmm0, (%rdi) -++; X64-NEXT: vmovw %xmm0, %eax -++; X64-NEXT: movb %ah, (%rdi) -+ ; X64-NEXT: retq -+ entry: -+ %i3 = shufflevector <8 x i16> %arg, <8 x i16> , <8 x i32> -+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr46532.ll b/llvm/test/CodeGen/X86/pr46532.ll -+--- a/llvm/test/CodeGen/X86/pr46532.ll -++++ b/llvm/test/CodeGen/X86/pr46532.ll -+@@ -7,6 +7,8 @@ -+ ; CHECK-NEXT: movq (%rax), %rax -+ ; CHECK-NEXT: vxorps %xmm0, %xmm0, %xmm0 -+ ; CHECK-NEXT: vmovaps %xmm0, 32(%rax) -++; CHECK-NEXT: vmovsd {{.*#+}} xmm0 = [4294967295,4294967295,0,0] -++; CHECK-NEXT: vmaskmovps %ymm0, %ymm0, (%rax) -+ while.1.body.preheader: -+ %0 = load ptr, ptr undef, align 8, !invariant.load !0, !dereferenceable !1, !align !2 -+ %1 = getelementptr inbounds i8, ptr %0, i64 32 -+diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td -+--- a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td -++++ b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td -+@@ -1,71 +0,0 @@ -+-// RUN: split-file %s %t -+-// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ -+-// RUN: | FileCheck %t/valid.td -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-bug.td 2>&1 \ -+-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=user-sgpr-init16-bug --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-priority.td 2>&1 \ -+-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=required-export-priority --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-timing.td 2>&1 \ -+-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=gfx125x-lowest-rate-wmma --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-banks.td 2>&1 \ -+-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=ldsbankcount32 --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-lds-size.td 2>&1 \ -+-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=addressablelocalmemorysize65536 --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/capability.td 2>&1 \ -+-// RUN: | FileCheck %t/capability.td -DFILE=%t/capability.td --implicit-check-not="error:" -+- -+-//--- common.td -+-include "AMDGPU.td" -+- -+-def EmptyGPU : AMDGPUProcessorModel<"gfx999", NoSchedModel, [], [9, 9, 9]>; -+- -+-//--- valid.td -+-include "common.td" -+-// Check the real generic targets, including gfx11's combined workarounds, -+-// gfx12-5's fixup and WMMA timing, and gfx9-4's LDS size and bank count. -+-// CHECK: GK_GFX9_4_GENERIC -+-// CHECK: GK_GFX11_GENERIC -+-// CHECK: GK_GFX12_5_GENERIC -+- -+-//--- missing-bug.td -+-include "common.td" -+-defvar CheckedFeature = FeatureUserSGPRInit16Bug; -+-include "missing-generic.td" -+- -+-//--- missing-priority.td -+-include "common.td" -+-defvar CheckedFeature = FeatureRequiredExportPriority; -+-include "missing-generic.td" -+- -+-//--- missing-timing.td -+-include "common.td" -+-defvar CheckedFeature = FeatureGFX125xLowestRateWMMA; -+-include "missing-generic.td" -+- -+-//--- missing-banks.td -+-include "common.td" -+-defvar CheckedFeature = FeatureLDSBankCount32; -+-include "missing-generic.td" -+- -+-//--- missing-lds-size.td -+-include "common.td" -+-defvar CheckedFeature = FeatureAddressableLocalMemorySize65536; -+-include "missing-generic.td" -+- -+-//--- missing-generic.td -+-// Classified backend-only features still need support from a covered GPU. -+-// Included files use the native path separator after the include directory. -+-// CHECK: [[DIR]]{{[/\\]}}missing-generic.td:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature '[[FEATURE]]' not supported by any covered GPU -+-def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [CheckedFeature], [9, 9, 0]> { -+- let CoveredGPUs = [EmptyGPU]; -+-} -+- -+-//--- capability.td -+-include "common.td" -+-def WMMAGPU : AMDGPUProcessorModel<"gfx998", NoSchedModel, -+- [FeatureWMMA256bInsts], [9, 9, 8]>; -+-// Instruction support keeps the default all-covered-GPUs rule. -+-// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature 'wmma-256b-insts' not supported by covered GPU 'gfx999' -+-def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [FeatureWMMA256bInsts], [9, 9, 0]> { -+- let CoveredGPUs = [WMMAGPU, EmptyGPU]; -+-} -+diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td -+--- a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td -++++ b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td -+@@ -1,33 +1,66 @@ -+ // RUN: split-file %s %t -+-// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ -++// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/valid.td \ -+ // RUN: | FileCheck %t/valid.td -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ -+ // RUN: | FileCheck %t/missing-any.td -DFILE=%t/missing-any.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ -+ // RUN: | FileCheck %t/missing-all.td -DFILE=%t/missing-all.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ -+ // RUN: | FileCheck %t/implied-all.td -DFILE=%t/implied-all.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ -+ // RUN: | FileCheck %t/missing-private.td -DFILE=%t/missing-private.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-private.td 2>&1 \ -++// RUN: | FileCheck %t/implied-private.td -DFILE=%t/implied-private.td --implicit-check-not="error:" -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ -+ // RUN: | FileCheck %t/numeric-mismatch.td -DFILE=%t/numeric-mismatch.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ -+ // RUN: | FileCheck %t/empty-members.td -DFILE=%t/empty-members.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ -+ // RUN: | FileCheck %t/self-member.td -DFILE=%t/self-member.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ -+ // RUN: | FileCheck %t/generic-member.td -DFILE=%t/generic-member.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ -+ // RUN: | FileCheck %t/pseudo-member.td -DFILE=%t/pseudo-member.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ -+ // RUN: | FileCheck %t/unregistered-member.td -DFILE=%t/unregistered-member.td --implicit-check-not="error:" -+-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ -++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ -+ // RUN: | FileCheck %t/r600-member.td -DFILE=%t/r600-member.td --implicit-check-not="error:" - -- // ALL: declare noundef nonnull ptr @_Znaj( -+ //--- common.td -+-include "llvm/Target/Target.td" -++// Exercise the validation rules with small synthetic targets. The real targets -++// are validated when AMDGPUTargetParserDef.inc is generated during the build. -++// Stand-ins for the Target.td classes used by AMDGPUTargetParser.td, limited to -++// the fields the backend reads. -++class InlineBehavior; -++def InlineIgnore : InlineBehavior; -++def InlineDefault : InlineBehavior; -++class SubtargetFeature i = [], -++ InlineBehavior ib = InlineDefault> { -++ string Name = n; -++ string FieldName = f; -++ string Value = v; -++ string Desc = d; -++ list Implies = i; -++ InlineBehavior InlineBehavior = ib; -++} -++ -++class SchedMachineModel; -++def NoSchedModel : SchedMachineModel; -++class ProcessorItineraries; -++def NoItineraries : ProcessorItineraries; + } +- +-+//--- implied-private.td +-+include "common.td" +-+def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [], [9, 0, 0]>; +-+// An unclassified backend-only parent must not hide a classified dependency. +-+// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-generic' exposes feature 'private' not supported by any covered GPU +-+def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, [FeatureImpliesPrivate], [9, 0, 0]> { +-+ let CoveredGPUs = [GFX900]; +-+} + --+void *f3(unsigned long N) { --+ // SANE: call noalias noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] --+ // SANENOT: call noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] --+ return ::operator new(N); -++class Processor f> { -++ string Name = n; -++ ProcessorItineraries ProcItin = pi; -++ list Features = f; +- //--- numeric-mismatch.td +- include "common.td" +- def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [FeatureAnySmall], [9, 0, 0]>; +-diff -ruN --strip-trailing-cr a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp +---- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp +-+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp +-@@ -14,11 +14,12 @@ +- using namespace mlir; +- using namespace mlir::tosa; +- +--// Building this ~5,000-line generated initializer map under MemorySanitizer +--// leads to extreme compile times (>3.5 minutes) during InstCombine and greedy +--// register allocation due to shadow/origin instrumentation overhead. Disable +--// optimization under MSan to keep build times manageable. +--#if LLVM_MEMORY_SANITIZER_BUILD +-+// Building this ~5,000-line generated initializer map under HWAddressSanitizer +-+// or MemorySanitizer leads to extreme compile times (>3.5 minutes) during +-+// InstCombine and greedy register allocation due to instrumentation +-+// overhead. Disable optimization under MSan/HWAsan to keep build times +-+// manageable. +-+#if LLVM_MEMORY_SANITIZER_BUILD || LLVM_HWADDRESS_SANITIZER_BUILD +- __attribute__((optnone)) +- #else +- LLVM_ATTRIBUTE_MINSIZE +-diff -ruN --strip-trailing-cr a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel +---- a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel +-+++ b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel +-@@ -66,7 +66,7 @@ +- +- bool_flag( +- name = "llvm_enable_lzma", +-- build_setting_default = False, +-+ build_setting_default = True, +- ) +- +- config_setting( +-@@ -74,11 +74,7 @@ +- flag_values = {":llvm_enable_lzma": "true"}, +- ) +- +--# Unlike zlib and zstd there is no liblzma module in the Bazel registry, so this +--# links the system library rather than building one from source. That is why it +--# is a plain cc_library with linkopts instead of a cc_library_wrapper, and why +--# it defaults to off. +--cc_library( +-+cc_library_wrapper( +- name = "lzma", +- defines = select({ +- ":llvm_lzma_enabled": [ +-@@ -88,9 +84,9 @@ +- "LLVM_ENABLE_LZMA=0", +- ], +- }), +-- linkopts = select({ +-+ deps = select({ +- ":llvm_lzma_enabled": [ +-- "-llzma", +-+ "@llvm_xz//:lzma", +- ], +- "//conditions:default": [], +- }), +-diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel b/utils/bazel/MODULE.bazel +---- a/utils/bazel/MODULE.bazel +-+++ b/utils/bazel/MODULE.bazel +-@@ -29,6 +29,7 @@ +- bazel_dep(name = "nanobind", version = "2.12.0.bcr.1") +- bazel_dep(name = "libpfm", version = "4.13.0", repo_name = "pfm") +- bazel_dep(name = "vulkan_headers", version = "1.4.349") +-+bazel_dep(name = "xz", version = "5.4.5.bcr.9", repo_name = "llvm_xz") +- +- bazel_dep(name = "emsdk", version = "6.0.2", dev_dependency = True) +- bazel_dep(name = "llvm", version = "0.8.20", dev_dependency = True) +-diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel.lock b/utils/bazel/MODULE.bazel.lock +---- a/utils/bazel/MODULE.bazel.lock +-+++ b/utils/bazel/MODULE.bazel.lock +-@@ -55,7 +55,8 @@ +- "https://bcr.bazel.build/modules/bazel_lib/3.0.0/MODULE.bazel": "22b70b80ac89ad3f3772526cd9feee2fa412c2b01933fea7ed13238a448d370d", +- "https://bcr.bazel.build/modules/bazel_lib/3.1.0/MODULE.bazel": "6809765c14e3c766a9b9286c7b0ec56ed87a73326e48fe01749f0c0fdcfe3287", +- "https://bcr.bazel.build/modules/bazel_lib/3.2.2/MODULE.bazel": "e2c890c8a515d6bca9c66d47718aa9e44b458fde64ec7204b8030bf2d349058c", +-- "https://bcr.bazel.build/modules/bazel_lib/3.2.2/source.json": "9e84e115c20e14652c5c21401ae85ff4daa8702e265b5c0b3bf89353f17aa212", +-+ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/MODULE.bazel": "cbc920fff00a55b2fd663d40850560a6fb17fd5ef3a9f3d6cbf0d32908dd8a40", +-+ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/source.json": "7a9c2b06a113c4962d329bb68befdb90d0447788a19aabaa32b476a0458e7d28", +- "https://bcr.bazel.build/modules/bazel_skylib/1.0.3/MODULE.bazel": "bcb0fd896384802d1ad283b4e4eb4d718eebd8cb820b0a2c3a347fb971afd9d8", +- "https://bcr.bazel.build/modules/bazel_skylib/1.1.1/MODULE.bazel": "1add3e7d93ff2e6998f9e118022c84d163917d912f5afafb3058e3d2f1545b5e", +- "https://bcr.bazel.build/modules/bazel_skylib/1.2.0/MODULE.bazel": "44fe84260e454ed94ad326352a698422dbe372b21a1ac9f3eab76eb531223686", +-@@ -313,6 +314,8 @@ +- "https://bcr.bazel.build/modules/windows_support/0.4.1/source.json": "bdb72e6717b00135e0a4a8512b3f6b626df0baea0bde74bf7312c7484f791287", +- "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/MODULE.bazel": "b573395fe63aef4299ba095173e2f62ccfee5ad9bbf7acaa95dba73af9fc2b38", +- "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/source.json": "3f3fbaeafecaf629877ad152a2c9def21f8d330d91aa94c5dc75bbb98c10b8b8", +-+ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/MODULE.bazel": "ff7cf8a5334a96b7b9ffcbe73418d45929ce125efebac39c1ee7d14bd2ae7a9b", +-+ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/source.json": "abf61f3c1630eab7b81722f710060b6f31e7c160511bcbfaad74bb4385883051", +- "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/MODULE.bazel": "86130924115bf8dfc4aba286fb49a2ec8d32430a21771db9ccf5381acf0d10cc", +- "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/source.json": "7aa4112eb2fbb99757e97cf0b7a4e2bdb3464b2b3cfb351bc807ffcb34bd349f", +- "https://bcr.bazel.build/modules/zlib-ng/2.3.3/MODULE.bazel": "0afa781a5f354b4fd811a5c9a086e111daf50e09d8e98ab1cf7c46eeacc33aac", +++define void @store_low1_nxv8f16_splat(ptr %0) { +++; CHECK-LABEL: store_low1_nxv8f16_splat: +++; CHECK: // %bb.0: +++; CHECK-NEXT: fmov h0, #1.00000000 +++; CHECK-NEXT: str h0, [x0] +++; CHECK-NEXT: ret +++ %2 = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 1) +++ tail call void @llvm.masked.store.nxv8f16.p0( splat (half 0xH3C00), ptr align 2 %0, %2) +++ ret void ++} -++class ProcessorModel f> -++ : Processor { -++ SchedMachineModel SchedModel = m; - +} - + -- // SANE: attributes [[ATTR]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } --+// SANE: attributes [[ATTR2]] = { allocsize(0) } -- // SANENOT: attributes [[ATTR]] = { builtin allocsize(0) } --+// SANENOT: attributes [[ATTR2]] = { allocsize(0) } --diff -ruN --strip-trailing-cr a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp ----- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp --+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp --@@ -16145,18 +16145,21 @@ -+ include "AMDGPUTargetParser.td" -+-def MyTarget : Target; - -- InstructionCost VecCost = VectorCost(CommonCost); -- // Check if the current node must be resized, if the parent node is not --- // resized. --+ // resized. Nodes without a parent (roots of the splat gather --+ // subtrees) are resized back to their original type when the reusing --+ // gather nodes are emitted. -- if (It != MinBWs.end() && !UnaryInstruction::isCast(E->getOpcode()) && -- E->Idx != 0 && -- (E->getOpcode() != Instruction::Load || E->UserTreeIndex)) { -- const EdgeInfo &EI = E->UserTreeIndex; --- if (!EI.UserTE->hasState() || --+ if (!EI.UserTE || !EI.UserTE->hasState() || -- EI.UserTE->getOpcode() != Instruction::Select || -- EI.EdgeIdx != 0) { -- auto UserBWIt = MinBWs.find(EI.UserTE); -- Type *UserScalarTy = --- (EI.UserTE->isGather() || --- EI.UserTE->State == TreeEntry::SplitVectorize) --+ !EI.UserTE ? OrigScalarTy --+ : (EI.UserTE->isGather() || --+ EI.UserTE->State == TreeEntry::SplitVectorize) -- ? EI.UserTE->Scalars.front()->getType() -- : EI.UserTE->getOperand(EI.EdgeIdx).front()->getType(); -- if (UserBWIt != MinBWs.end()) --diff -ruN --strip-trailing-cr a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll ----- a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll --+++ b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll --@@ -0,0 +1,75 @@ --+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py --+; RUN: opt -passes=slp-vectorizer -S -mtriple=aarch64-linux-gnu < %s | FileCheck %s --+ --+define double @test() { --+; CHECK-LABEL: @test( --+; CHECK-NEXT: bbl: --+; CHECK-NEXT: [[SELECT:%.*]] = select i1 false, i64 0, i64 0 --+; CHECK-NEXT: [[AND:%.*]] = and i64 [[SELECT]], 0 --+; CHECK-NEXT: [[ICMP:%.*]] = icmp eq i64 [[AND]], 0 --+; CHECK-NEXT: [[TMP0:%.*]] = insertelement <2 x i1> poison, i1 true, i64 0 --+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x i1> [[TMP0]], <2 x i1> poison, <2 x i32> zeroinitializer --+; CHECK-NEXT: [[TMP2:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer --+; CHECK-NEXT: [[TMP3:%.*]] = bitcast <2 x double> [[TMP2]] to <2 x i64> --+; CHECK-NEXT: [[TMP4:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer --+; CHECK-NEXT: [[TMP5:%.*]] = bitcast <2 x double> [[TMP4]] to <2 x i64> --+; CHECK-NEXT: [[TMP6:%.*]] = insertelement <2 x i1> poison, i1 [[ICMP]], i64 0 --+; CHECK-NEXT: [[TMP7:%.*]] = shufflevector <2 x i1> [[TMP6]], <2 x i1> poison, <2 x i32> zeroinitializer --+; CHECK-NEXT: [[TMP8:%.*]] = select <2 x i1> [[TMP7]], <2 x double> zeroinitializer, <2 x double> zeroinitializer --+; CHECK-NEXT: [[TMP9:%.*]] = bitcast <2 x double> [[TMP8]] to <2 x i64> --+; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x i64> poison, i64 [[SELECT]], i64 0 --+; CHECK-NEXT: [[TMP11:%.*]] = shufflevector <2 x i64> [[TMP10]], <2 x i64> poison, <2 x i32> zeroinitializer --+; CHECK-NEXT: [[TMP12:%.*]] = xor <2 x i64> [[TMP11]], [[TMP9]] --+; CHECK-NEXT: [[TMP13:%.*]] = bitcast <2 x i64> [[TMP12]] to <2 x double> --+; CHECK-NEXT: [[TMP14:%.*]] = insertelement <2 x i64> poison, i64 0, i64 0 --+; CHECK-NEXT: [[TMP15:%.*]] = shufflevector <2 x i64> [[TMP14]], <2 x i64> poison, <2 x i32> zeroinitializer --+; CHECK-NEXT: [[TMP16:%.*]] = xor <2 x i64> [[TMP15]], [[TMP3]] --+; CHECK-NEXT: [[TMP17:%.*]] = bitcast <2 x i64> [[TMP16]] to <2 x double> --+; CHECK-NEXT: [[TMP18:%.*]] = xor <2 x i64> [[TMP15]], [[TMP5]] --+; CHECK-NEXT: [[TMP19:%.*]] = fmul <2 x double> [[TMP13]], [[TMP17]] --+; CHECK-NEXT: [[TMP20:%.*]] = bitcast <2 x i64> [[TMP18]] to <2 x double> --+; CHECK-NEXT: [[TMP21:%.*]] = fmul <2 x double> [[TMP19]], [[TMP20]] --+; CHECK-NEXT: [[TMP22:%.*]] = extractelement <2 x double> [[TMP21]], i64 0 --+; CHECK-NEXT: [[TMP23:%.*]] = extractelement <2 x double> [[TMP21]], i64 1 --+; CHECK-NEXT: [[FADD:%.*]] = fadd double [[TMP22]], [[TMP23]] --+; CHECK-NEXT: ret double [[FADD]] --+; --+bbl: --+ %select = select i1 false, i64 0, i64 0 --+ %and = and i64 %select, 0 --+ %icmp = icmp eq i64 %and, 0 --+ %select1 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 --+ %bitcast = bitcast double %select1 to i64 --+ %xor = xor i64 %select, %bitcast --+ %bitcast2 = bitcast i64 %xor to double --+ %select4 = select i1 false, i64 0, i64 0 --+ %icmp5 = icmp eq i64 %select4, 0 --+ %select6 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 --+ %bitcast7 = bitcast double %select6 to i64 --+ %xor8 = xor i64 %select4, %bitcast7 --+ %bitcast9 = bitcast i64 %xor8 to double --+ %select10 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 --+ %bitcast11 = bitcast double %select10 to i64 --+ %select12 = select i1 false, i64 0, i64 0 --+ %icmp13 = icmp eq i64 %select12, 0 --+ %select14 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 --+ %bitcast15 = bitcast double %select14 to i64 --+ %xor16 = xor i64 %select12, %bitcast15 --+ %select17 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 --+ %bitcast18 = bitcast double %select17 to i64 --+ %select3 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 --+ %bitcast19 = bitcast double %select3 to i64 --+ %xor20 = xor i64 %select, %bitcast19 --+ %bitcast21 = bitcast i64 %xor20 to double --+ %xor22 = xor i64 %select4, %bitcast11 --+ %bitcast23 = bitcast i64 %xor22 to double --+ %fmul = fmul double %bitcast21, %bitcast23 --+ %xor24 = xor i64 %select12, %bitcast18 --+ %bitcast25 = bitcast i64 %xor24 to double --+ %fmul26 = fmul double %fmul, %bitcast25 --+ %fmul27 = fmul double %bitcast2, %bitcast9 --+ %bitcast28 = bitcast i64 %xor16 to double --+ %fmul29 = fmul double %fmul27, %bitcast28 --+ %fadd = fadd double %fmul26, %fmul29 --+ ret double %fadd -+ def FeatureAll : SubtargetFeature<"all", "HasAll", "true", "">; -+ def FeatureAnyA : AMDGPUGenericAnyFeature<"any-a", "HasAnyA", "true", "">; -+@@ -37,7 +70,13 @@ -+ def FeatureAnySmall : NumericFeature<4>; -+ def FeatureAnyB : NumericFeature<8>; -+ def FeatureAnyLarge : NumericFeature<16>; -+-def FeaturePrivate : AMDGPUGenericAnyFeature<"private", "HasPrivate", "true", "">; -++// Classification inherited through a subclass also applies to backend-only -++// features, which must participate in validation without being emitted. -++class PrivateFeature : AMDGPUGenericAnyFeature< -++ "private", "HasPrivate", "true", "">; -++def FeaturePrivate : PrivateFeature; -++def FeatureImpliesPrivate : SubtargetFeature< -++ "implies-private", "HasImpliesPrivate", "true", "", [FeaturePrivate]>; -+ def FeatureImpliesA : SubtargetFeature< -+ "implies-a", "HasImpliesA", "true", "", [FeatureAnyA]>; -+ def FeatureImpliesB : SubtargetFeature< -+@@ -55,7 +94,7 @@ -+ def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, -+ [FeatureAll, FeatureAnyA, FeatureAnySmall], [9, 0, 0]>; -+ def GFX901 : AMDGPUProcessorModel<"gfx901", NoSchedModel, -+- [FeatureAll, FeatureImpliesB, FeaturePrivate], [9, 0, 1]>; -++ [FeatureAll, FeatureImpliesB, FeatureImpliesPrivate], [9, 0, 1]>; -+ // Different members supply different marked features. Follow Implies on both -+ // the generic and the members, and check marked backend-only features too. -+ def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, -+@@ -103,6 +142,15 @@ -+ let CoveredGPUs = [GFX900]; -+ } -+ -++//--- implied-private.td -++include "common.td" -++def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [], [9, 0, 0]>; -++// An unclassified backend-only parent must not hide a classified dependency. -++// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-generic' exposes feature 'private' not supported by any covered GPU -++def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, [FeatureImpliesPrivate], [9, 0, 0]> { -++ let CoveredGPUs = [GFX900]; - +} -++ -+ //--- numeric-mismatch.td -+ include "common.td" -+ def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [FeatureAnySmall], [9, 0, 0]>; -+diff -ruN --strip-trailing-cr a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp -+--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp -++++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp -+@@ -14,11 +14,12 @@ -+ using namespace mlir; -+ using namespace mlir::tosa; -+ -+-// Building this ~5,000-line generated initializer map under MemorySanitizer -+-// leads to extreme compile times (>3.5 minutes) during InstCombine and greedy -+-// register allocation due to shadow/origin instrumentation overhead. Disable -+-// optimization under MSan to keep build times manageable. -+-#if LLVM_MEMORY_SANITIZER_BUILD -++// Building this ~5,000-line generated initializer map under HWAddressSanitizer -++// or MemorySanitizer leads to extreme compile times (>3.5 minutes) during -++// InstCombine and greedy register allocation due to instrumentation -++// overhead. Disable optimization under MSan/HWAsan to keep build times -++// manageable. -++#if LLVM_MEMORY_SANITIZER_BUILD || LLVM_HWADDRESS_SANITIZER_BUILD -+ __attribute__((optnone)) -+ #else -+ LLVM_ATTRIBUTE_MINSIZE -+diff -ruN --strip-trailing-cr a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel -+--- a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel -++++ b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel -+@@ -66,7 +66,7 @@ -+ -+ bool_flag( -+ name = "llvm_enable_lzma", -+- build_setting_default = False, -++ build_setting_default = True, -+ ) -+ -+ config_setting( -+@@ -74,11 +74,7 @@ -+ flag_values = {":llvm_enable_lzma": "true"}, -+ ) -+ -+-# Unlike zlib and zstd there is no liblzma module in the Bazel registry, so this -+-# links the system library rather than building one from source. That is why it -+-# is a plain cc_library with linkopts instead of a cc_library_wrapper, and why -+-# it defaults to off. -+-cc_library( -++cc_library_wrapper( -+ name = "lzma", -+ defines = select({ -+ ":llvm_lzma_enabled": [ -+@@ -88,9 +84,9 @@ -+ "LLVM_ENABLE_LZMA=0", -+ ], -+ }), -+- linkopts = select({ -++ deps = select({ -+ ":llvm_lzma_enabled": [ -+- "-llzma", -++ "@llvm_xz//:lzma", -+ ], -+ "//conditions:default": [], -+ }), -+diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel b/utils/bazel/MODULE.bazel -+--- a/utils/bazel/MODULE.bazel -++++ b/utils/bazel/MODULE.bazel -+@@ -29,6 +29,7 @@ -+ bazel_dep(name = "nanobind", version = "2.12.0.bcr.1") -+ bazel_dep(name = "libpfm", version = "4.13.0", repo_name = "pfm") -+ bazel_dep(name = "vulkan_headers", version = "1.4.349") -++bazel_dep(name = "xz", version = "5.4.5.bcr.9", repo_name = "llvm_xz") -+ -+ bazel_dep(name = "emsdk", version = "6.0.2", dev_dependency = True) -+ bazel_dep(name = "llvm", version = "0.8.20", dev_dependency = True) -+diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel.lock b/utils/bazel/MODULE.bazel.lock -+--- a/utils/bazel/MODULE.bazel.lock -++++ b/utils/bazel/MODULE.bazel.lock -+@@ -55,7 +55,8 @@ -+ "https://bcr.bazel.build/modules/bazel_lib/3.0.0/MODULE.bazel": "22b70b80ac89ad3f3772526cd9feee2fa412c2b01933fea7ed13238a448d370d", -+ "https://bcr.bazel.build/modules/bazel_lib/3.1.0/MODULE.bazel": "6809765c14e3c766a9b9286c7b0ec56ed87a73326e48fe01749f0c0fdcfe3287", -+ "https://bcr.bazel.build/modules/bazel_lib/3.2.2/MODULE.bazel": "e2c890c8a515d6bca9c66d47718aa9e44b458fde64ec7204b8030bf2d349058c", -+- "https://bcr.bazel.build/modules/bazel_lib/3.2.2/source.json": "9e84e115c20e14652c5c21401ae85ff4daa8702e265b5c0b3bf89353f17aa212", -++ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/MODULE.bazel": "cbc920fff00a55b2fd663d40850560a6fb17fd5ef3a9f3d6cbf0d32908dd8a40", -++ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/source.json": "7a9c2b06a113c4962d329bb68befdb90d0447788a19aabaa32b476a0458e7d28", -+ "https://bcr.bazel.build/modules/bazel_skylib/1.0.3/MODULE.bazel": "bcb0fd896384802d1ad283b4e4eb4d718eebd8cb820b0a2c3a347fb971afd9d8", -+ "https://bcr.bazel.build/modules/bazel_skylib/1.1.1/MODULE.bazel": "1add3e7d93ff2e6998f9e118022c84d163917d912f5afafb3058e3d2f1545b5e", -+ "https://bcr.bazel.build/modules/bazel_skylib/1.2.0/MODULE.bazel": "44fe84260e454ed94ad326352a698422dbe372b21a1ac9f3eab76eb531223686", -+@@ -313,6 +314,8 @@ -+ "https://bcr.bazel.build/modules/windows_support/0.4.1/source.json": "bdb72e6717b00135e0a4a8512b3f6b626df0baea0bde74bf7312c7484f791287", -+ "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/MODULE.bazel": "b573395fe63aef4299ba095173e2f62ccfee5ad9bbf7acaa95dba73af9fc2b38", -+ "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/source.json": "3f3fbaeafecaf629877ad152a2c9def21f8d330d91aa94c5dc75bbb98c10b8b8", -++ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/MODULE.bazel": "ff7cf8a5334a96b7b9ffcbe73418d45929ce125efebac39c1ee7d14bd2ae7a9b", -++ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/source.json": "abf61f3c1630eab7b81722f710060b6f31e7c160511bcbfaad74bb4385883051", -+ "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/MODULE.bazel": "86130924115bf8dfc4aba286fb49a2ec8d32430a21771db9ccf5381acf0d10cc", -+ "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/source.json": "7aa4112eb2fbb99757e97cf0b7a4e2bdb3464b2b3cfb351bc807ffcb34bd349f", -+ "https://bcr.bazel.build/modules/zlib-ng/2.3.3/MODULE.bazel": "0afa781a5f354b4fd811a5c9a086e111daf50e09d8e98ab1cf7c46eeacc33aac", diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl -index d929da09..88052c28 100644 +index 88052c28..ea77ef6b 100644 --- a/third_party/llvm/workspace.bzl +++ b/third_party/llvm/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") def repo(name): """Imports LLVM.""" -- LLVM_COMMIT = "a22291abb68744cb0b70978eb506a7c5592a3309" -- LLVM_SHA256 = "cf19b82f0dfbef1e62afe520ba4ee6de4f4823524d7804490e88d815f822c8e9" -+ LLVM_COMMIT = "e56c2cefc3e7978de4a2d799fa3abec77d75af21" -+ LLVM_SHA256 = "affcad31c9455e95faca9d4de38e4a49b73c23bbf48b405d264ff10a6bfd3baf" +- LLVM_COMMIT = "e56c2cefc3e7978de4a2d799fa3abec77d75af21" +- LLVM_SHA256 = "affcad31c9455e95faca9d4de38e4a49b73c23bbf48b405d264ff10a6bfd3baf" ++ LLVM_COMMIT = "6ec87e7ec4b3ca6f665a02e05d0e5fe9c37cd6cc" ++ LLVM_SHA256 = "75300dbf9232d135c1ee2819216b7b7597fd9d36a90312c2b0c45938097d806f" tf_http_archive( name = name, -diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch -index fb55a869..81f3bf2f 100644 ---- a/third_party/stablehlo/temporary.patch -+++ b/third_party/stablehlo/temporary.patch -@@ -123,6 +123,328 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU - - /// Extracts integer values from the attribute |elements|. - SmallVector extract1DVector(DenseIntElementsAttr elements); -+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir -+--- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir -++++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir -+@@ -38,7 +38,7 @@ -+ -+ // CHECK-LABEL: @concatenate -+ func.func @concatenate(%arg0 : tensor<3x3xf32>, %arg1 : tensor<3x3xf32>) -> tensor<6x3xf32> { -+- // CHECK: tosa.concat %arg0, %arg1 {axis = 0 : i32} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> -++ // CHECK: tosa.concat %arg0, %arg1 axis(0) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> -+ %0 = "stablehlo.concatenate"(%arg0, %arg1) {dimension = 0 : i64} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> -+ return %0 : tensor<6x3xf32> -+ } -+@@ -61,10 +61,10 @@ -+ -+ // CHECK-LABEL: @dot_vector_vector -+ func.func @dot_vector_vector(%arg0 : tensor<3xf32>, %arg1 : tensor<3xf32>) -> tensor { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -74,10 +74,10 @@ -+ -+ // CHECK-LABEL: @dot_vector_matrix -+ func.func @dot_vector_matrix(%arg0 : tensor<2xf32>, %arg1 : tensor<2x3xf32>) -> tensor<3xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -87,10 +87,10 @@ -+ -+ // CHECK-LABEL: @dot_matrix_vector -+ func.func @dot_matrix_vector(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3xf32>) -> tensor<2xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -100,10 +100,10 @@ -+ -+ // CHECK-LABEL: @dot_matrix_matrix -+ func.func @dot_matrix_matrix(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3x4xf32>) -> tensor<2x4xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -113,10 +113,10 @@ -+ -+ // CHECK-LABEL: @dot_general_vector_vector -+ func.func @dot_general_vector_vector(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -126,10 +126,10 @@ -+ -+ // CHECK-LABEL: @dot_general_vector_matrix -+ func.func @dot_general_vector_matrix(%arg0: tensor<2xf32>, %arg1: tensor<2x3xf32>) -> tensor<3xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -139,10 +139,10 @@ -+ -+ // CHECK-LABEL: @dot_general_matrix_vector -+ func.func @dot_general_matrix_vector(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+@@ -152,10 +152,10 @@ -+ -+ // CHECK-LABEL: @dot_general_matrix_matrix -+ func.func @dot_general_matrix_matrix(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] -+- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3> -++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3> -+ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] -+ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] -+ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] -+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir -+--- stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir -++++ stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir -+@@ -11,10 +11,10 @@ -+ // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.add %[[V0]], %[[V1]] : tensor<2x2xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> -+ %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) -+ -> tensor<2x2x!quant.uniform> -+@@ -32,10 +32,10 @@ -+ // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.subtract %[[V0]], %[[V1]] : tensor<2x2xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> -+ %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) -+ -> tensor<2x2x!quant.uniform> -+@@ -52,10 +52,10 @@ -+ // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1717986918> : tensor<1xi32>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.multiply %[[V0]], %[[V1]] : tensor<2x2xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> -+ %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) -+ -> tensor<2x2x!quant.uniform> -+@@ -74,10 +74,10 @@ -+ // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.divide %[[V0]], %[[V1]] : tensor<2x2xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> -+ %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) -+ -> tensor<2x2x!quant.uniform> -+@@ -97,10 +97,10 @@ -+ // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.maximum %[[V0]], %[[V1]] : tensor<2x2xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> -+ %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) -+ -> tensor<2x2x!quant.uniform> -+@@ -120,10 +120,10 @@ -+ // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.minimum %[[V0]], %[[V1]] : tensor<2x2xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> -+ %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) -+ -> tensor<2x2x!quant.uniform> -+@@ -140,9 +140,9 @@ -+ // CHECK-DAG: %[[SHIFT30:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V1:.+]] = stablehlo.abs %[[V0]] : tensor<20x20xi32> -+- // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: return %[[V3]] : tensor<20x20x!quant.uniform> -+ %0 = "stablehlo.abs"(%arg0) : (tensor<20x20x!quant.uniform>) -> tensor<20x20x!quant.uniform> -+ return %0 : tensor<20x20x!quant.uniform> -+@@ -159,8 +159,8 @@ -+ // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.compare GE, %[[V0]], %[[V1]], TOTALORDER : -+ // CHECK: return %[[V2]] -+ %0 = stablehlo.compare GE, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform>, tensor<20x20x!quant.uniform>) -> tensor<20x20xi1> -+@@ -177,8 +177,8 @@ -+ // CHECK-DAG: %[[SHIFT15:.+]] = "tosa.const"() <{values = dense<15> : tensor<1xi8>}> -+ // CHECK-DAG: %[[ZP16_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi16>}> -+ // CHECK-DAG: %[[ZP32_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> -+- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -+- // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} -++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -++ // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) -+ // CHECK: %[[V2:.+]] = stablehlo.compare LT, %[[V0]], %[[V1]], TOTALORDER : -+ // CHECK: return %[[V2]] -+ %0 = stablehlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform>, tensor<20x20x!quant.uniform>) -> tensor<20x20xi1> -+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir -+--- stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir -++++ stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir -+@@ -7,7 +7,7 @@ -+ %shift = "tosa.const"() {values = dense<13> : tensor<1xi8>} : () -> tensor<1xi8> -+ %input_zp = "tosa.const"() {values = dense<-1> : tensor<1xi8>} : () -> tensor<1xi8> -+ %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> -+- %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} : -++ %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) : -+ (tensor<2x2x!quant.uniform>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x2xi32> -+ -+ // convert input quantized type to storage type -+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir -+--- stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir -++++ stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir -+@@ -18,7 +18,7 @@ -+ func.func @iota_dimension_0() -> tensor<4x8xf32> { -+ // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() -+ // CHECK-SAME{LITERAL}: <{values = dense<[[0.000000e+00], [1.000000e+00], [2.000000e+00], [3.000000e+00]]> : tensor<4x1xf32>}> : () -> tensor<4x1xf32> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 8]> : vector<2xindex>} : () -> !tosa.shape<2> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 8]> : vector<2xindex>) : () -> !tosa.shape<2> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]] -+ %0 = "stablehlo.iota"() {iota_dimension = 0 : i64} : () -> (tensor<4x8xf32>) -+ return %0 : tensor<4x8xf32> -+@@ -28,7 +28,7 @@ -+ func.func @iota_dimension_1() -> tensor<4x8xi32> { -+ // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() -+ // CHECK-SAME{LITERAL}: <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>}> : () -> tensor<1x8xi32> -+- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[4, 1]> : vector<2xindex>} : () -> !tosa.shape<2> -++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[4, 1]> : vector<2xindex>) : () -> !tosa.shape<2> -+ // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]] -+ %0 = "stablehlo.iota"() {iota_dimension = 1 : i64} : () -> (tensor<4x8xi32>) -+ return %0 : tensor<4x8xi32> -+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir -+--- stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir -++++ stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir -+@@ -2,7 +2,7 @@ -+ -+ // CHECK-LABEL: @concatenate -+ func.func @concatenate(%arg0 : tensor<5x2xf32>, %arg1 : tensor<5x5xf32>, %arg2 : tensor<5x7xf32>) -> tensor<5x14xf32> { -+- // CHECK: tosa.concat %arg0, %arg1, %arg2 {axis = 1 : i32} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> -++ // CHECK: tosa.concat %arg0, %arg1, %arg2 axis(1) : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> -+ %0 = "stablehlo.concatenate"(%arg0, %arg1, %arg2) {dimension = 1 : i64} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> -+ return %0 : tensor<5x14xf32> -+ } -+diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir -+--- stablehlo/stablehlo/conversions/tosa/tests/unary.mlir -++++ stablehlo/stablehlo/conversions/tosa/tests/unary.mlir -+@@ -79,8 +79,8 @@ -+ -+ // CHECK-LABEL: @slice -+ func.func @slice(%arg : tensor<4x3xf32>) -> tensor<2x2xf32> { -+- // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape {values = dense<[2, 1]> : tensor<2xindex>} : () -> !tosa.shape<2> -+- // CHECK-DAG: %[[START:.*]] = tosa.const_shape {values = dense<2> : tensor<2xindex>} : () -> !tosa.shape<2> -++ // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape values(dense<[2, 1]> : tensor<2xindex>) : () -> !tosa.shape<2> -++ // CHECK-DAG: %[[START:.*]] = tosa.const_shape values(dense<2> : tensor<2xindex>) : () -> !tosa.shape<2> -+ // CHECK: tosa.slice %arg0, %[[SIZE]], %[[START]] -+ %0 = "stablehlo.slice"(%arg) { -+ start_indices = array, -+@@ -123,7 +123,7 @@ -+ -+ // CHECK-LABEL: @transpose -+ func.func @transpose(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { -+- // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 {perms = array} -++ // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 perms([2, 1, 0]) -+ %0 = "stablehlo.transpose"(%arg0) {permutation = array} : (tensor<1x2x3xf32>) -> tensor<3x2x1xf32> -+ return %0 : tensor<3x2x1xf32> -+ } -+@@ -158,7 +158,7 @@ -+ -+ // CHECK-LABEL: @reshape -+ func.func @reshape(%arg0 : tensor<2x3xf32>) -> tensor<6xf32> { -+- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<6> : tensor<1xindex>} : () -> !tosa.shape<1> -++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<6> : tensor<1xindex>) : () -> !tosa.shape<1> -+ // CHECK: tosa.reshape %arg0, %[[VAR0]] -+ %0 = "stablehlo.reshape"(%arg0) : (tensor<2x3xf32>) -> tensor<6xf32> -+ return %0 : tensor<6xf32> - diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp - --- stablehlo/stablehlo/dialect/ChloBytecode.cpp - +++ stablehlo/stablehlo/dialect/ChloBytecode.cpp diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl index 2e45bac52ee8ca..cf564a294d7d8a 100644 --- a/third_party/xla/third_party/shardy/workspace.bzl +++ b/third_party/xla/third_party/shardy/workspace.bzl @@ -18,8 +18,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): - SHARDY_COMMIT = "dc6f5f05c411667e7f2f6a2d3552276c6f4b1a43" - SHARDY_SHA256 = "0130b0bfa450f1fb2a44da6264aec45e81607896790a799d5fc75c6ff08bafa2" + SHARDY_COMMIT = "88e76b98bc30632838472fe80f84caf2c2b18693" + SHARDY_SHA256 = "25df7b5104c269505e4aec5ff70b5ed434b58deed0d8aebe1276625a9d481827" tf_http_archive( name = "shardy", diff --git a/third_party/xla/xla/backends/autotuner/config_runner.cc b/third_party/xla/xla/backends/autotuner/config_runner.cc index 3855e17ac5d7cf..cac73b30b38457 100644 --- a/third_party/xla/xla/backends/autotuner/config_runner.cc +++ b/third_party/xla/xla/backends/autotuner/config_runner.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" @@ -130,11 +131,23 @@ ConfigRunner::ConfigProfile ConfigRunner::ProfileCandidate( profiler_->Profile(candidate.executable.get(), input_buffers); if (!profile_result.ok()) { + if (options_.enable_correctness_check) { + absl::Status redzone = profiler_->CheckInputBuffers(input_buffers); + if (!redzone.ok()) { + return ConfigProfile{ + /*config=*/std::move(candidate.config), + /*failure=*/ + Failure{FailureKind::kRedzoneCheckFailed, redzone.ToString()}, + }; + } + } + FailureKind failure_kind = + absl::StrContains(profile_result.status().message(), "Redzone") + ? FailureKind::kRedzoneCheckFailed + : FailureKind::kExecutionFailed; return ConfigProfile{ /*config=*/std::move(candidate.config), - /*failure=*/ - Failure{FailureKind::kExecutionFailed, - profile_result.status().ToString()}, + /*failure=*/Failure{failure_kind, profile_result.status().ToString()}, }; } diff --git a/third_party/xla/xla/backends/autotuner/config_runner_test.cc b/third_party/xla/xla/backends/autotuner/config_runner_test.cc index 116a99a44f63d0..8c7e0bd7eb5feb 100644 --- a/third_party/xla/xla/backends/autotuner/config_runner_test.cc +++ b/third_party/xla/xla/backends/autotuner/config_runner_test.cc @@ -187,6 +187,72 @@ TEST_F(ConfigRunnerTest, ProfileAllRedzoneCheckFailed) { ConfigRunner::FailureKind::kRedzoneCheckFailed); } +TEST_F(ConfigRunnerTest, ProfileAllWarmupRedzoneCheckFailed) { + auto profiler = std::make_unique(); + EXPECT_CALL(*profiler, CreateInputBuffers(_, _)).WillOnce([] { + return std::make_unique(); + }); + EXPECT_CALL(*profiler, Profile(_, _)) + .WillOnce(Return(absl::InternalError( + "Autotuning candidate rejected: kernel wrote past its allocated " + "buffer. Redzone mismatch in LHS redzone of buffer 0x30fd08000"))); + EXPECT_CALL(*profiler, CheckInputBuffers(_)) + .WillOnce(Return(absl::OkStatus())); + + ConfigRunner::CorrectnessCheckOptions options; + options.enable_correctness_check = true; + ASSERT_OK_AND_ASSIGN(auto runner, + ConfigRunner::Create(std::move(profiler), options)); + + MockCodegenBackend backend; + std::vector candidates; + candidates.push_back({ + /*config=*/{&backend, GetTestConfig("test_config_1")}, + /*executable=*/std::make_unique(), + }); + + ASSERT_OK_AND_ASSIGN(auto profiles, + runner->ProfileAll(std::move(candidates))); + ASSERT_THAT(profiles, SizeIs(1)); + ASSERT_TRUE(profiles[0].failure.has_value()); + EXPECT_EQ(profiles[0].failure->kind, + ConfigRunner::FailureKind::kRedzoneCheckFailed); + EXPECT_EQ(profiles[0].ToProto().failure().kind(), + AutotuneResult::REDZONE_MODIFIED); + EXPECT_EQ(profiles[0].ToFailedConfigsProto().kind(), + autotuner::FailedConfigs::REDZONE_CHECK_FAILED); +} + +TEST_F(ConfigRunnerTest, ProfileAllExecutionErrorWithInputRedzoneModified) { + auto profiler = std::make_unique(); + EXPECT_CALL(*profiler, CreateInputBuffers(_, _)).WillOnce([] { + return std::make_unique(); + }); + EXPECT_CALL(*profiler, Profile(_, _)) + .WillOnce(Return(absl::InternalError("execution crash"))); + EXPECT_CALL(*profiler, CheckInputBuffers(_)) + .WillOnce(Return(absl::InternalError("input redzone modified"))); + + ConfigRunner::CorrectnessCheckOptions options; + options.enable_correctness_check = true; + ASSERT_OK_AND_ASSIGN(auto runner, + ConfigRunner::Create(std::move(profiler), options)); + + MockCodegenBackend backend; + std::vector candidates; + candidates.push_back({ + /*config=*/{&backend, GetTestConfig("test_config_1")}, + /*executable=*/std::make_unique(), + }); + + ASSERT_OK_AND_ASSIGN(auto profiles, + runner->ProfileAll(std::move(candidates))); + ASSERT_THAT(profiles, SizeIs(1)); + ASSERT_TRUE(profiles[0].failure.has_value()); + EXPECT_EQ(profiles[0].failure->kind, + ConfigRunner::FailureKind::kRedzoneCheckFailed); +} + TEST_F(ConfigRunnerTest, FiltersWrongResultsAgainstTrusted) { ConfigRunner::CorrectnessCheckOptions options; options.enable_correctness_check = true; diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl b/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl index 2c590d2b9e4ef1..c1d7076f67cddc 100644 --- a/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl @@ -47,6 +47,14 @@ REGULAR_HLO_FILES = [ "layer_norm_1x4096x768_bf16.hlo", "mean_axis_1x4096x1024_bf16.hlo", "mha_block_1x12x128x64_bf16.hlo", + "pairwise_biot_savart_4096_f64.hlo", + "pairwise_biot_savart_grad_2048_f64.hlo", + "pairwise_gravity_4096_f64.hlo", + "pairwise_gravity_grad_2048_f64.hlo", + "pairwise_min_distance_grad_2048_f64.hlo", + "pairwise_rbf_gradient_4096_f32.hlo", + "pairwise_rbf_gradient_grad_2048_f32.hlo", + "pairwise_sum_distances_grad_2048_f32.hlo", "resnet50_keras_jax_f32.hlo", "sort_full_1024x4096_bf16.hlo", "sum_axis_1x4096x1024_bf16.hlo", diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_4096_f64.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_4096_f64.hlo new file mode 100644 index 00000000000000..7eb39c597c574d --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_4096_f64.hlo @@ -0,0 +1,105 @@ +HloModule jit_biot_savart, entry_computation_layout={(f64[4096,3]{1,0}, f64[4096,3]{1,0}, f64[4096]{0})->f64[4096,3]{1,0}} + +region_0.1 { + reduce_sum.3 = f64[] parameter(0) + reduce_sum.4 = f64[] parameter(1) + ROOT reduce_sum.5 = f64[] add(reduce_sum.3, reduce_sum.4) +} + +norm.2 { + Arg_0.1 = f64[4096,4096,3]{2,1,0} parameter(0) + mul.1 = f64[4096,4096,3]{2,1,0} multiply(Arg_0.1, Arg_0.1) + constant.2 = f64[] constant(0) + reduce_sum.7 = f64[4096,4096]{1,0} reduce(mul.1, constant.2), dimensions={2}, to_apply=region_0.1 + broadcast_in_dim.3 = f64[4096,4096,1]{2,1,0} reshape(reduce_sum.7) + ROOT sqrt.1 = f64[4096,4096,1]{2,1,0} sqrt(broadcast_in_dim.3) +} + +cross.3 { + Arg_0.3 = f64[4096,3]{1,0} parameter(0) + slice.7 = f64[4096,1]{1,0} slice(Arg_0.3), slice={[0:4096], [1:2]} + broadcast_in_dim_squeeze.6 = f64[1,4096]{1,0} reshape(slice.7) + mul.26 = f64[1,4096]{1,0} broadcast(broadcast_in_dim_squeeze.6), dimensions={0,1} + mul.27 = f64[4096]{0} reshape(mul.26) + mul.28 = f64[4096,4096]{1,0} broadcast(mul.27), dimensions={1} + Arg_1.1 = f64[4096,4096,3]{2,1,0} parameter(1) + slice.11 = f64[4096,4096,1]{2,1,0} slice(Arg_1.1), slice={[0:4096], [0:4096], [2:3]} + squeeze.5 = f64[4096,4096]{1,0} reshape(slice.11) + mul.29 = f64[4096,4096]{1,0} multiply(mul.28, squeeze.5) + slice.8 = f64[4096,1]{1,0} slice(Arg_0.3), slice={[0:4096], [2:3]} + broadcast_in_dim_squeeze.7 = f64[1,4096]{1,0} reshape(slice.8) + mul.30 = f64[1,4096]{1,0} broadcast(broadcast_in_dim_squeeze.7), dimensions={0,1} + mul.31 = f64[4096]{0} reshape(mul.30) + mul.32 = f64[4096,4096]{1,0} broadcast(mul.31), dimensions={1} + slice.10 = f64[4096,4096,1]{2,1,0} slice(Arg_1.1), slice={[0:4096], [0:4096], [1:2]} + squeeze.4 = f64[4096,4096]{1,0} reshape(slice.10) + mul.33 = f64[4096,4096]{1,0} multiply(mul.32, squeeze.4) + sub.10 = f64[4096,4096]{1,0} subtract(mul.29, mul.33) + broadcast_in_dim.7 = f64[1,4096,4096]{2,1,0} reshape(sub.10) + broadcast_in_dim_squeeze.8 = f64[1,4096]{1,0} reshape(slice.8) + mul.34 = f64[1,4096]{1,0} broadcast(broadcast_in_dim_squeeze.8), dimensions={0,1} + mul.35 = f64[4096]{0} reshape(mul.34) + mul.36 = f64[4096,4096]{1,0} broadcast(mul.35), dimensions={1} + slice.9 = f64[4096,4096,1]{2,1,0} slice(Arg_1.1), slice={[0:4096], [0:4096], [0:1]} + squeeze.3 = f64[4096,4096]{1,0} reshape(slice.9) + mul.37 = f64[4096,4096]{1,0} multiply(mul.36, squeeze.3) + slice.6 = f64[4096,1]{1,0} slice(Arg_0.3), slice={[0:4096], [0:1]} + broadcast_in_dim_squeeze.9 = f64[1,4096]{1,0} reshape(slice.6) + mul.38 = f64[1,4096]{1,0} broadcast(broadcast_in_dim_squeeze.9), dimensions={0,1} + mul.39 = f64[4096]{0} reshape(mul.38) + mul.40 = f64[4096,4096]{1,0} broadcast(mul.39), dimensions={1} + mul.41 = f64[4096,4096]{1,0} multiply(mul.40, squeeze.5) + sub.11 = f64[4096,4096]{1,0} subtract(mul.37, mul.41) + broadcast_in_dim.8 = f64[1,4096,4096]{2,1,0} reshape(sub.11) + broadcast_in_dim_squeeze.10 = f64[1,4096]{1,0} reshape(slice.6) + mul.42 = f64[1,4096]{1,0} broadcast(broadcast_in_dim_squeeze.10), dimensions={0,1} + mul.43 = f64[4096]{0} reshape(mul.42) + mul.44 = f64[4096,4096]{1,0} broadcast(mul.43), dimensions={1} + mul.45 = f64[4096,4096]{1,0} multiply(mul.44, squeeze.4) + broadcast_in_dim_squeeze.11 = f64[1,4096]{1,0} reshape(slice.7) + mul.46 = f64[1,4096]{1,0} broadcast(broadcast_in_dim_squeeze.11), dimensions={0,1} + mul.47 = f64[4096]{0} reshape(mul.46) + mul.48 = f64[4096,4096]{1,0} broadcast(mul.47), dimensions={1} + mul.49 = f64[4096,4096]{1,0} multiply(mul.48, squeeze.3) + sub.12 = f64[4096,4096]{1,0} subtract(mul.45, mul.49) + broadcast_in_dim.9 = f64[1,4096,4096]{2,1,0} reshape(sub.12) + concatenate.1 = f64[3,4096,4096]{2,1,0} concatenate(broadcast_in_dim.7, broadcast_in_dim.8, broadcast_in_dim.9), dimensions={0} + ROOT transpose.1 = f64[4096,4096,3]{1,0,2} transpose(concatenate.1), dimensions={1,2,0} +} + +region_1.4 { + reduce_sum.11 = f64[] parameter(0) + reduce_sum.12 = f64[] parameter(1) + ROOT reduce_sum.13 = f64[] add(reduce_sum.11, reduce_sum.12) +} + +ENTRY main.5 { + w.1 = f64[4096]{0} parameter(2) + stack.4 = f64[4096,1]{1,0} reshape(w.1) + stack.5 = f64[4096,1]{1,0} reshape(w.1) + stack.6 = f64[4096,1]{1,0} reshape(w.1) + stack.7 = f64[4096,3]{1,0} concatenate(stack.4, stack.5, stack.6), dimensions={1} + pos.1 = f64[4096,3]{1,0} parameter(0) + broadcast_in_dim.10 = f64[4096,1,3]{2,1,0} reshape(pos.1) + sub.13 = f64[4096,1,3]{2,1,0} broadcast(broadcast_in_dim.10), dimensions={0,1,2} + sub.14 = f64[4096,3]{1,0} reshape(sub.13) + sub.15 = f64[4096,4096,3]{2,1,0} broadcast(sub.14), dimensions={0,2} + src.1 = f64[4096,3]{1,0} parameter(1) + broadcast_in_dim.11 = f64[1,4096,3]{2,1,0} reshape(src.1) + sub.16 = f64[1,4096,3]{2,1,0} broadcast(broadcast_in_dim.11), dimensions={0,1,2} + sub.17 = f64[4096,3]{1,0} reshape(sub.16) + sub.18 = f64[4096,4096,3]{2,1,0} broadcast(sub.17), dimensions={1,2} + sub.19 = f64[4096,4096,3]{2,1,0} subtract(sub.15, sub.18) + jit_cross_.1 = f64[4096,4096,3]{1,0,2} call(stack.7, sub.19), to_apply=cross.3 + jit_norm_.1 = f64[4096,4096,1]{2,1,0} call(sub.19), to_apply=norm.2 + integer_pow.2 = f64[4096,4096,1]{2,1,0} multiply(jit_norm_.1, jit_norm_.1) + integer_pow.3 = f64[4096,4096,1]{2,1,0} multiply(integer_pow.2, jit_norm_.1) + div.4 = f64[4096,4096,1]{2,1,0} broadcast(integer_pow.3), dimensions={0,1,2} + div.5 = f64[4096,4096]{1,0} reshape(div.4) + div.6 = f64[4096,4096,3]{2,1,0} broadcast(div.5), dimensions={0,1} + div.7 = f64[4096,4096,3]{1,0,2} divide(jit_cross_.1, div.6) + constant.3 = f64[] constant(0) + ROOT reduce_sum.15 = f64[4096,3]{1,0} reduce(div.7, constant.3), dimensions={1}, to_apply=region_1.4 +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_grad_2048_f64.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_grad_2048_f64.hlo new file mode 100644 index 00000000000000..2594759a2c99d2 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_biot_savart_grad_2048_f64.hlo @@ -0,0 +1,278 @@ +HloModule jit_loss, entry_computation_layout={(f64[2048,3]{1,0}, f64[2048,3]{1,0}, f64[2048]{0})->(f64[2048,3]{1,0}, f64[2048,3]{1,0})} + +region_0.1 { + reduce_sum.3 = f64[] parameter(0) + reduce_sum.4 = f64[] parameter(1) + ROOT reduce_sum.5 = f64[] add(reduce_sum.3, reduce_sum.4) +} + +norm.2 { + Arg_0.1 = f64[2048,2048,3]{2,1,0} parameter(0) + mul.3 = f64[2048,2048,3]{2,1,0} multiply(Arg_0.1, Arg_0.1) + constant.7 = f64[] constant(0) + reduce_sum.7 = f64[2048,2048]{1,0} reduce(mul.3, constant.7), dimensions={2}, to_apply=region_0.1 + broadcast_in_dim.4 = f64[2048,2048,1]{2,1,0} reshape(reduce_sum.7) + sqrt.1 = f64[2048,2048,1]{2,1,0} sqrt(broadcast_in_dim.4) + constant.6 = f64[] constant(0.5) + div.2 = f64[2048,2048,1]{2,1,0} broadcast(constant.6), dimensions={} + div.3 = f64[2048,2048,1]{2,1,0} divide(div.2, sqrt.1) + ROOT tuple.1 = (f64[2048,2048,1]{2,1,0}, f64[2048,2048,1]{2,1,0}) tuple(sqrt.1, div.3) +} + +cross.3 { + Arg_0.3 = f64[2048,3]{1,0} parameter(0) + slice.7 = f64[2048,1]{1,0} slice(Arg_0.3), slice={[0:2048], [1:2]} + broadcast_in_dim_squeeze.6 = f64[1,2048]{1,0} reshape(slice.7) + mul.28 = f64[1,2048]{1,0} broadcast(broadcast_in_dim_squeeze.6), dimensions={0,1} + mul.29 = f64[2048]{0} reshape(mul.28) + mul.30 = f64[2048,2048]{1,0} broadcast(mul.29), dimensions={1} + Arg_1.1 = f64[2048,2048,3]{2,1,0} parameter(1) + slice.11 = f64[2048,2048,1]{2,1,0} slice(Arg_1.1), slice={[0:2048], [0:2048], [2:3]} + squeeze.5 = f64[2048,2048]{1,0} reshape(slice.11) + mul.31 = f64[2048,2048]{1,0} multiply(mul.30, squeeze.5) + slice.8 = f64[2048,1]{1,0} slice(Arg_0.3), slice={[0:2048], [2:3]} + broadcast_in_dim_squeeze.7 = f64[1,2048]{1,0} reshape(slice.8) + mul.32 = f64[1,2048]{1,0} broadcast(broadcast_in_dim_squeeze.7), dimensions={0,1} + mul.33 = f64[2048]{0} reshape(mul.32) + mul.34 = f64[2048,2048]{1,0} broadcast(mul.33), dimensions={1} + slice.10 = f64[2048,2048,1]{2,1,0} slice(Arg_1.1), slice={[0:2048], [0:2048], [1:2]} + squeeze.4 = f64[2048,2048]{1,0} reshape(slice.10) + mul.35 = f64[2048,2048]{1,0} multiply(mul.34, squeeze.4) + sub.10 = f64[2048,2048]{1,0} subtract(mul.31, mul.35) + broadcast_in_dim.8 = f64[1,2048,2048]{2,1,0} reshape(sub.10) + broadcast_in_dim_squeeze.8 = f64[1,2048]{1,0} reshape(slice.8) + mul.36 = f64[1,2048]{1,0} broadcast(broadcast_in_dim_squeeze.8), dimensions={0,1} + mul.37 = f64[2048]{0} reshape(mul.36) + mul.38 = f64[2048,2048]{1,0} broadcast(mul.37), dimensions={1} + slice.9 = f64[2048,2048,1]{2,1,0} slice(Arg_1.1), slice={[0:2048], [0:2048], [0:1]} + squeeze.3 = f64[2048,2048]{1,0} reshape(slice.9) + mul.39 = f64[2048,2048]{1,0} multiply(mul.38, squeeze.3) + slice.6 = f64[2048,1]{1,0} slice(Arg_0.3), slice={[0:2048], [0:1]} + broadcast_in_dim_squeeze.9 = f64[1,2048]{1,0} reshape(slice.6) + mul.40 = f64[1,2048]{1,0} broadcast(broadcast_in_dim_squeeze.9), dimensions={0,1} + mul.41 = f64[2048]{0} reshape(mul.40) + mul.42 = f64[2048,2048]{1,0} broadcast(mul.41), dimensions={1} + mul.43 = f64[2048,2048]{1,0} multiply(mul.42, squeeze.5) + sub.11 = f64[2048,2048]{1,0} subtract(mul.39, mul.43) + broadcast_in_dim.9 = f64[1,2048,2048]{2,1,0} reshape(sub.11) + broadcast_in_dim_squeeze.10 = f64[1,2048]{1,0} reshape(slice.6) + mul.44 = f64[1,2048]{1,0} broadcast(broadcast_in_dim_squeeze.10), dimensions={0,1} + mul.45 = f64[2048]{0} reshape(mul.44) + mul.46 = f64[2048,2048]{1,0} broadcast(mul.45), dimensions={1} + mul.47 = f64[2048,2048]{1,0} multiply(mul.46, squeeze.4) + broadcast_in_dim_squeeze.11 = f64[1,2048]{1,0} reshape(slice.7) + mul.48 = f64[1,2048]{1,0} broadcast(broadcast_in_dim_squeeze.11), dimensions={0,1} + mul.49 = f64[2048]{0} reshape(mul.48) + mul.50 = f64[2048,2048]{1,0} broadcast(mul.49), dimensions={1} + mul.51 = f64[2048,2048]{1,0} multiply(mul.50, squeeze.3) + sub.12 = f64[2048,2048]{1,0} subtract(mul.47, mul.51) + broadcast_in_dim.10 = f64[1,2048,2048]{2,1,0} reshape(sub.12) + concatenate.1 = f64[3,2048,2048]{2,1,0} concatenate(broadcast_in_dim.8, broadcast_in_dim.9, broadcast_in_dim.10), dimensions={0} + transpose.1 = f64[2048,2048,3]{1,0,2} transpose(concatenate.1), dimensions={1,2,0} + ROOT tuple.3 = (f64[2048,2048,3]{1,0,2}, f64[1,2048]{1,0}, f64[1,2048]{1,0}, f64[1,2048]{1,0}, f64[1,2048]{1,0}, /*index=5*/f64[1,2048]{1,0}, f64[1,2048]{1,0}) tuple(transpose.1, broadcast_in_dim_squeeze.6, broadcast_in_dim_squeeze.7, broadcast_in_dim_squeeze.8, broadcast_in_dim_squeeze.9, broadcast_in_dim_squeeze.10, broadcast_in_dim_squeeze.11) +} + +region_1.4 { + reduce_sum.11 = f64[] parameter(0) + reduce_sum.12 = f64[] parameter(1) + ROOT reduce_sum.13 = f64[] add(reduce_sum.11, reduce_sum.12) +} + +region_2.5 { + reduce_sum.18 = f64[] parameter(0) + reduce_sum.19 = f64[] parameter(1) + ROOT reduce_sum.20 = f64[] add(reduce_sum.18, reduce_sum.19) +} + +region_3.6 { + reduce_sum.25 = f64[] parameter(0) + reduce_sum.26 = f64[] parameter(1) + ROOT reduce_sum.27 = f64[] add(reduce_sum.25, reduce_sum.26) +} + +region_4.7 { + reduce_sum.32 = f64[] parameter(0) + reduce_sum.33 = f64[] parameter(1) + ROOT reduce_sum.34 = f64[] add(reduce_sum.32, reduce_sum.33) +} + +region_5.8 { + reduce_sum.39 = f64[] parameter(0) + reduce_sum.40 = f64[] parameter(1) + ROOT reduce_sum.41 = f64[] add(reduce_sum.39, reduce_sum.40) +} + +cross_0.9 { + Arg_3.1 = f64[1,2048]{1,0} parameter(3) + mul.93 = f64[1,2048]{1,0} broadcast(Arg_3.1), dimensions={0,1} + mul.94 = f64[2048]{0} reshape(mul.93) + mul.95 = f64[2048,2048]{1,0} broadcast(mul.94), dimensions={1} + Arg_6.1 = f64[2048,2048,3]{2,1,0} parameter(6) + transpose.3 = f64[3,2048,2048]{0,2,1} transpose(Arg_6.1), dimensions={2,0,1} + split.4 = f64[1,2048,2048]{2,1,0} slice(transpose.3), slice={[1:2], [0:2048], [0:2048]} + constant.9 = f64[] constant(0) + reduce_sum.44 = f64[2048,2048]{1,0} reduce(split.4, constant.9), dimensions={0}, to_apply=region_4.7 + neg.5 = f64[2048,2048]{1,0} negate(reduce_sum.44) + mul.96 = f64[2048,2048]{1,0} multiply(mul.95, neg.5) + Arg_0.5 = f64[1,2048]{1,0} parameter(0) + mul.105 = f64[1,2048]{1,0} broadcast(Arg_0.5), dimensions={0,1} + mul.106 = f64[2048]{0} reshape(mul.105) + mul.107 = f64[2048,2048]{1,0} broadcast(mul.106), dimensions={1} + split.3 = f64[1,2048,2048]{2,1,0} slice(transpose.3), slice={[0:1], [0:2048], [0:2048]} + reduce_sum.45 = f64[2048,2048]{1,0} reduce(split.3, constant.9), dimensions={0}, to_apply=region_5.8 + mul.108 = f64[2048,2048]{1,0} multiply(mul.107, reduce_sum.45) + add_any.7 = f64[2048,2048]{1,0} add(mul.96, mul.108) + broadcast_in_dim.15 = f64[2048,2048,1]{2,1,0} reshape(add_any.7) + pad.3 = f64[2048,2048,3]{2,1,0} pad(broadcast_in_dim.15, constant.9), padding=0_0x0_0x2_0 + Arg_4.1 = f64[1,2048]{1,0} parameter(4) + mul.89 = f64[1,2048]{1,0} broadcast(Arg_4.1), dimensions={0,1} + mul.90 = f64[2048]{0} reshape(mul.89) + mul.91 = f64[2048,2048]{1,0} broadcast(mul.90), dimensions={1} + split.5 = f64[1,2048,2048]{2,1,0} slice(transpose.3), slice={[2:3], [0:2048], [0:2048]} + reduce_sum.43 = f64[2048,2048]{1,0} reduce(split.5, constant.9), dimensions={0}, to_apply=region_3.6 + mul.92 = f64[2048,2048]{1,0} multiply(mul.91, reduce_sum.43) + Arg_1.3 = f64[1,2048]{1,0} parameter(1) + mul.101 = f64[1,2048]{1,0} broadcast(Arg_1.3), dimensions={0,1} + mul.102 = f64[2048]{0} reshape(mul.101) + mul.103 = f64[2048,2048]{1,0} broadcast(mul.102), dimensions={1} + neg.6 = f64[2048,2048]{1,0} negate(reduce_sum.45) + mul.104 = f64[2048,2048]{1,0} multiply(mul.103, neg.6) + add_any.6 = f64[2048,2048]{1,0} add(mul.92, mul.104) + broadcast_in_dim.16 = f64[2048,2048,1]{2,1,0} reshape(add_any.6) + pad.4 = f64[2048,2048,3]{2,1,0} pad(broadcast_in_dim.16, constant.9), padding=0_0x0_0x1_1 + add_any.8 = f64[2048,2048,3]{2,1,0} add(pad.3, pad.4) + Arg_5.1 = f64[1,2048]{1,0} parameter(5) + mul.85 = f64[1,2048]{1,0} broadcast(Arg_5.1), dimensions={0,1} + mul.86 = f64[2048]{0} reshape(mul.85) + mul.87 = f64[2048,2048]{1,0} broadcast(mul.86), dimensions={1} + neg.4 = f64[2048,2048]{1,0} negate(reduce_sum.43) + mul.88 = f64[2048,2048]{1,0} multiply(mul.87, neg.4) + Arg_2.1 = f64[1,2048]{1,0} parameter(2) + mul.97 = f64[1,2048]{1,0} broadcast(Arg_2.1), dimensions={0,1} + mul.98 = f64[2048]{0} reshape(mul.97) + mul.99 = f64[2048,2048]{1,0} broadcast(mul.98), dimensions={1} + mul.100 = f64[2048,2048]{1,0} multiply(mul.99, reduce_sum.44) + add_any.5 = f64[2048,2048]{1,0} add(mul.88, mul.100) + broadcast_in_dim.17 = f64[2048,2048,1]{2,1,0} reshape(add_any.5) + pad.5 = f64[2048,2048,3]{2,1,0} pad(broadcast_in_dim.17, constant.9), padding=0_0x0_0x0_2 + ROOT add_any.9 = f64[2048,2048,3]{2,1,0} add(add_any.8, pad.5) +} + +region_6.10 { + reduce_sum.49 = f64[] parameter(0) + reduce_sum.50 = f64[] parameter(1) + ROOT reduce_sum.51 = f64[] add(reduce_sum.49, reduce_sum.50) +} + +norm_1.11 { + Arg_0.7 = f64[2048,2048,3]{2,1,0} parameter(0) + Arg_2.3 = f64[2048,2048,1]{2,1,0} parameter(2) + Arg_1.5 = f64[2048,2048,1]{2,1,0} parameter(1) + mul.112 = f64[2048,2048,1]{2,1,0} multiply(Arg_2.3, Arg_1.5) + constant.11 = f64[] constant(0) + reduce_sum.53 = f64[2048,2048]{1,0} reduce(mul.112, constant.11), dimensions={2}, to_apply=region_6.10 + broadcast_in_dim.19 = f64[2048,2048,3]{2,1,0} broadcast(reduce_sum.53), dimensions={0,1} + mul.113 = f64[2048,2048,3]{2,1,0} multiply(Arg_0.7, broadcast_in_dim.19) + mul.114 = f64[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.19, Arg_0.7) + ROOT add_any.11 = f64[2048,2048,3]{2,1,0} add(mul.113, mul.114) +} + +region_7.12 { + reduce_sum.57 = f64[] parameter(0) + reduce_sum.58 = f64[] parameter(1) + ROOT reduce_sum.59 = f64[] add(reduce_sum.57, reduce_sum.58) +} + +region_8.13 { + reduce_sum.64 = f64[] parameter(0) + reduce_sum.65 = f64[] parameter(1) + ROOT reduce_sum.66 = f64[] add(reduce_sum.64, reduce_sum.65) +} + +region_9.14 { + reduce_sum.71 = f64[] parameter(0) + reduce_sum.72 = f64[] parameter(1) + ROOT reduce_sum.73 = f64[] add(reduce_sum.71, reduce_sum.72) +} + +region_10.15 { + reduce_sum.78 = f64[] parameter(0) + reduce_sum.79 = f64[] parameter(1) + ROOT reduce_sum.80 = f64[] add(reduce_sum.78, reduce_sum.79) +} + +ENTRY main.16 { + a_2_.1 = f64[2048]{0} parameter(2) + stack.4 = f64[2048,1]{1,0} reshape(a_2_.1) + stack.5 = f64[2048,1]{1,0} reshape(a_2_.1) + stack.6 = f64[2048,1]{1,0} reshape(a_2_.1) + stack.7 = f64[2048,3]{1,0} concatenate(stack.4, stack.5, stack.6), dimensions={1} + a_0_.1 = f64[2048,3]{1,0} parameter(0) + broadcast_in_dim.21 = f64[2048,1,3]{2,1,0} reshape(a_0_.1) + sub.13 = f64[2048,1,3]{2,1,0} broadcast(broadcast_in_dim.21), dimensions={0,1,2} + sub.14 = f64[2048,3]{1,0} reshape(sub.13) + sub.15 = f64[2048,2048,3]{2,1,0} broadcast(sub.14), dimensions={0,2} + a_1_.1 = f64[2048,3]{1,0} parameter(1) + broadcast_in_dim.22 = f64[1,2048,3]{2,1,0} reshape(a_1_.1) + sub.16 = f64[1,2048,3]{2,1,0} broadcast(broadcast_in_dim.22), dimensions={0,1,2} + sub.17 = f64[2048,3]{1,0} reshape(sub.16) + sub.18 = f64[2048,2048,3]{2,1,0} broadcast(sub.17), dimensions={1,2} + sub.19 = f64[2048,2048,3]{2,1,0} subtract(sub.15, sub.18) + jvp_jit_cross__.8 = (f64[2048,2048,3]{1,0,2}, f64[1,2048]{1,0}, f64[1,2048]{1,0}, f64[1,2048]{1,0}, f64[1,2048]{1,0}, /*index=5*/f64[1,2048]{1,0}, f64[1,2048]{1,0}) call(stack.7, sub.19), to_apply=cross.3 + jvp_jit_cross__.10 = f64[1,2048]{1,0} get-tuple-element(jvp_jit_cross__.8), index=1 + jvp_jit_cross__.11 = f64[1,2048]{1,0} get-tuple-element(jvp_jit_cross__.8), index=2 + jvp_jit_cross__.12 = f64[1,2048]{1,0} get-tuple-element(jvp_jit_cross__.8), index=3 + jvp_jit_cross__.13 = f64[1,2048]{1,0} get-tuple-element(jvp_jit_cross__.8), index=4 + jvp_jit_cross__.14 = f64[1,2048]{1,0} get-tuple-element(jvp_jit_cross__.8), index=5 + jvp_jit_cross__.15 = f64[1,2048]{1,0} get-tuple-element(jvp_jit_cross__.8), index=6 + jvp_jit_cross__.9 = f64[2048,2048,3]{1,0,2} get-tuple-element(jvp_jit_cross__.8), index=0 + jvp_jit_norm__.3 = (f64[2048,2048,1]{2,1,0}, f64[2048,2048,1]{2,1,0}) call(sub.19), to_apply=norm.2 + jvp_jit_norm__.4 = f64[2048,2048,1]{2,1,0} get-tuple-element(jvp_jit_norm__.3), index=0 + integer_pow.3 = f64[2048,2048,1]{2,1,0} multiply(jvp_jit_norm__.4, jvp_jit_norm__.4) + integer_pow.4 = f64[2048,2048,1]{2,1,0} multiply(integer_pow.3, jvp_jit_norm__.4) + div.13 = f64[2048,2048,1]{2,1,0} broadcast(integer_pow.4), dimensions={0,1,2} + div.14 = f64[2048,2048]{1,0} reshape(div.13) + div.15 = f64[2048,2048,3]{2,1,0} broadcast(div.14), dimensions={0,1} + div.16 = f64[2048,2048,3]{1,0,2} divide(jvp_jit_cross__.9, div.15) + constant.15 = f64[] constant(0) + reduce_sum.82 = f64[2048,3]{1,0} reduce(div.16, constant.15), dimensions={1}, to_apply=region_1.4 + constant.12 = f64[] constant(2) + mul.115 = f64[2048,3]{1,0} broadcast(constant.12), dimensions={} + mul.119 = f64[2048,3]{1,0} multiply(reduce_sum.82, mul.115) + broadcast_in_dim.23 = f64[2048,2048,3]{2,1,0} broadcast(mul.119), dimensions={0,2} + div.18 = f64[2048,2048,1]{2,1,0} broadcast(integer_pow.4), dimensions={0,1,2} + div.19 = f64[2048,2048]{1,0} reshape(div.18) + div.20 = f64[2048,2048,3]{2,1,0} broadcast(div.19), dimensions={0,1} + div.21 = f64[2048,2048,3]{2,1,0} divide(broadcast_in_dim.23, div.20) + transpose_jvp_jit_cross___.1 = f64[2048,2048,3]{2,1,0} call(jvp_jit_cross__.10, jvp_jit_cross__.11, jvp_jit_cross__.12, jvp_jit_cross__.13, jvp_jit_cross__.14, jvp_jit_cross__.15, div.21), to_apply=cross_0.9 + jvp_jit_norm__.5 = f64[2048,2048,1]{2,1,0} get-tuple-element(jvp_jit_norm__.3), index=1 + constant.13 = f64[] constant(1) + broadcast_in_dim.20 = f64[2048,2048,1]{2,1,0} broadcast(constant.13), dimensions={} + mul.118 = f64[2048,2048,1]{2,1,0} multiply(integer_pow.4, integer_pow.4) + div.17 = f64[2048,2048,1]{2,1,0} divide(broadcast_in_dim.20, mul.118) + mul.120 = f64[2048,2048,1]{2,1,0} broadcast(div.17), dimensions={0,1,2} + mul.121 = f64[2048,2048]{1,0} reshape(mul.120) + mul.122 = f64[2048,2048,3]{2,1,0} broadcast(mul.121), dimensions={0,1} + mul.123 = f64[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.23, mul.122) + mul.124 = f64[2048,2048,3]{2,1,0} multiply(mul.123, jvp_jit_cross__.9) + reduce_sum.83 = f64[2048,2048]{1,0} reduce(mul.124, constant.15), dimensions={2}, to_apply=region_2.5 + reshape.3 = f64[2048,2048,1]{2,1,0} reshape(reduce_sum.83) + neg.8 = f64[2048,2048,1]{2,1,0} negate(reshape.3) + integer_pow.5 = f64[2048,2048,1]{2,1,0} multiply(jvp_jit_norm__.4, jvp_jit_norm__.4) + constant.14 = f64[] constant(3) + mul.116 = f64[2048,2048,1]{2,1,0} broadcast(constant.14), dimensions={} + mul.117 = f64[2048,2048,1]{2,1,0} multiply(integer_pow.5, mul.116) + mul.125 = f64[2048,2048,1]{2,1,0} multiply(neg.8, mul.117) + transpose_jvp_jit_norm___.1 = f64[2048,2048,3]{2,1,0} call(sub.19, jvp_jit_norm__.5, mul.125), to_apply=norm_1.11 + add_any.13 = f64[2048,2048,3]{2,1,0} add(transpose_jvp_jit_cross___.1, transpose_jvp_jit_norm___.1) + reduce_sum.84 = f64[2048,3]{1,0} reduce(add_any.13, constant.15), dimensions={1}, to_apply=region_7.12 + reshape.4 = f64[2048,1,3]{2,1,0} reshape(reduce_sum.84) + reduce_sum.87 = f64[2048,3]{1,0} reduce(reshape.4, constant.15), dimensions={1}, to_apply=region_10.15 + neg.9 = f64[2048,2048,3]{2,1,0} negate(add_any.13) + reduce_sum.85 = f64[2048,3]{1,0} reduce(neg.9, constant.15), dimensions={0}, to_apply=region_8.13 + reshape.5 = f64[1,2048,3]{2,1,0} reshape(reduce_sum.85) + reduce_sum.86 = f64[2048,3]{1,0} reduce(reshape.5, constant.15), dimensions={0}, to_apply=region_9.14 + ROOT tuple.5 = (f64[2048,3]{1,0}, f64[2048,3]{1,0}) tuple(reduce_sum.87, reduce_sum.86) +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_4096_f64.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_4096_f64.hlo new file mode 100644 index 00000000000000..c82b4e67bfd5e1 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_4096_f64.hlo @@ -0,0 +1,53 @@ +HloModule jit_gravity_accel, entry_computation_layout={(f64[4096,3]{1,0}, f64[4096,3]{1,0}, f64[4096]{0})->f64[4096,3]{1,0}} + +region_0.1 { + reduce_sum.3 = f64[] parameter(0) + reduce_sum.4 = f64[] parameter(1) + ROOT reduce_sum.5 = f64[] add(reduce_sum.3, reduce_sum.4) +} + +norm.2 { + Arg_0.1 = f64[4096,4096,3]{2,1,0} parameter(0) + mul.1 = f64[4096,4096,3]{2,1,0} multiply(Arg_0.1, Arg_0.1) + constant.2 = f64[] constant(0) + reduce_sum.7 = f64[4096,4096]{1,0} reduce(mul.1, constant.2), dimensions={2}, to_apply=region_0.1 + broadcast_in_dim.3 = f64[4096,4096,1]{2,1,0} reshape(reduce_sum.7) + ROOT sqrt.1 = f64[4096,4096,1]{2,1,0} sqrt(broadcast_in_dim.3) +} + +region_1.3 { + reduce_sum.11 = f64[] parameter(0) + reduce_sum.12 = f64[] parameter(1) + ROOT reduce_sum.13 = f64[] add(reduce_sum.11, reduce_sum.12) +} + +ENTRY main.4 { + w.1 = f64[4096]{0} parameter(2) + broadcast_in_dim.7 = f64[1,4096,1]{2,1,0} reshape(w.1) + mul.6 = f64[1,4096,1]{2,1,0} broadcast(broadcast_in_dim.7), dimensions={0,1,2} + mul.7 = f64[4096]{0} reshape(mul.6) + mul.8 = f64[4096,4096,3]{2,1,0} broadcast(mul.7), dimensions={1} + pos.1 = f64[4096,3]{1,0} parameter(0) + broadcast_in_dim.5 = f64[4096,1,3]{2,1,0} reshape(pos.1) + sub.7 = f64[4096,1,3]{2,1,0} broadcast(broadcast_in_dim.5), dimensions={0,1,2} + sub.8 = f64[4096,3]{1,0} reshape(sub.7) + sub.9 = f64[4096,4096,3]{2,1,0} broadcast(sub.8), dimensions={0,2} + src.1 = f64[4096,3]{1,0} parameter(1) + broadcast_in_dim.6 = f64[1,4096,3]{2,1,0} reshape(src.1) + sub.10 = f64[1,4096,3]{2,1,0} broadcast(broadcast_in_dim.6), dimensions={0,1,2} + sub.11 = f64[4096,3]{1,0} reshape(sub.10) + sub.12 = f64[4096,4096,3]{2,1,0} broadcast(sub.11), dimensions={1,2} + sub.13 = f64[4096,4096,3]{2,1,0} subtract(sub.9, sub.12) + mul.9 = f64[4096,4096,3]{2,1,0} multiply(mul.8, sub.13) + jit_norm_.1 = f64[4096,4096,1]{2,1,0} call(sub.13), to_apply=norm.2 + integer_pow.2 = f64[4096,4096,1]{2,1,0} multiply(jit_norm_.1, jit_norm_.1) + integer_pow.3 = f64[4096,4096,1]{2,1,0} multiply(integer_pow.2, jit_norm_.1) + div.4 = f64[4096,4096,1]{2,1,0} broadcast(integer_pow.3), dimensions={0,1,2} + div.5 = f64[4096,4096]{1,0} reshape(div.4) + div.6 = f64[4096,4096,3]{2,1,0} broadcast(div.5), dimensions={0,1} + div.7 = f64[4096,4096,3]{2,1,0} divide(mul.9, div.6) + constant.3 = f64[] constant(0) + ROOT reduce_sum.15 = f64[4096,3]{1,0} reduce(div.7, constant.3), dimensions={1}, to_apply=region_1.3 +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_grad_2048_f64.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_grad_2048_f64.hlo new file mode 100644 index 00000000000000..5c7868bf97104c --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_gravity_grad_2048_f64.hlo @@ -0,0 +1,147 @@ +HloModule jit_loss, entry_computation_layout={(f64[2048,3]{1,0}, f64[2048,3]{1,0}, f64[2048]{0})->(f64[2048,3]{1,0}, f64[2048,3]{1,0})} + +region_0.1 { + reduce_sum.3 = f64[] parameter(0) + reduce_sum.4 = f64[] parameter(1) + ROOT reduce_sum.5 = f64[] add(reduce_sum.3, reduce_sum.4) +} + +norm.2 { + Arg_0.1 = f64[2048,2048,3]{2,1,0} parameter(0) + mul.3 = f64[2048,2048,3]{2,1,0} multiply(Arg_0.1, Arg_0.1) + constant.7 = f64[] constant(0) + reduce_sum.7 = f64[2048,2048]{1,0} reduce(mul.3, constant.7), dimensions={2}, to_apply=region_0.1 + broadcast_in_dim.4 = f64[2048,2048,1]{2,1,0} reshape(reduce_sum.7) + sqrt.1 = f64[2048,2048,1]{2,1,0} sqrt(broadcast_in_dim.4) + constant.6 = f64[] constant(0.5) + div.2 = f64[2048,2048,1]{2,1,0} broadcast(constant.6), dimensions={} + div.3 = f64[2048,2048,1]{2,1,0} divide(div.2, sqrt.1) + ROOT tuple.1 = (f64[2048,2048,1]{2,1,0}, f64[2048,2048,1]{2,1,0}) tuple(sqrt.1, div.3) +} + +region_1.3 { + reduce_sum.11 = f64[] parameter(0) + reduce_sum.12 = f64[] parameter(1) + ROOT reduce_sum.13 = f64[] add(reduce_sum.11, reduce_sum.12) +} + +region_2.4 { + reduce_sum.18 = f64[] parameter(0) + reduce_sum.19 = f64[] parameter(1) + ROOT reduce_sum.20 = f64[] add(reduce_sum.18, reduce_sum.19) +} + +region_3.5 { + reduce_sum.25 = f64[] parameter(0) + reduce_sum.26 = f64[] parameter(1) + ROOT reduce_sum.27 = f64[] add(reduce_sum.25, reduce_sum.26) +} + +norm_0.6 { + Arg_0.3 = f64[2048,2048,3]{2,1,0} parameter(0) + Arg_2.1 = f64[2048,2048,1]{2,1,0} parameter(2) + Arg_1.1 = f64[2048,2048,1]{2,1,0} parameter(1) + mul.24 = f64[2048,2048,1]{2,1,0} multiply(Arg_2.1, Arg_1.1) + constant.9 = f64[] constant(0) + reduce_sum.29 = f64[2048,2048]{1,0} reduce(mul.24, constant.9), dimensions={2}, to_apply=region_3.5 + broadcast_in_dim.8 = f64[2048,2048,3]{2,1,0} broadcast(reduce_sum.29), dimensions={0,1} + mul.25 = f64[2048,2048,3]{2,1,0} multiply(Arg_0.3, broadcast_in_dim.8) + mul.26 = f64[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.8, Arg_0.3) + ROOT add_any.1 = f64[2048,2048,3]{2,1,0} add(mul.25, mul.26) +} + +region_4.7 { + reduce_sum.33 = f64[] parameter(0) + reduce_sum.34 = f64[] parameter(1) + ROOT reduce_sum.35 = f64[] add(reduce_sum.33, reduce_sum.34) +} + +region_5.8 { + reduce_sum.40 = f64[] parameter(0) + reduce_sum.41 = f64[] parameter(1) + ROOT reduce_sum.42 = f64[] add(reduce_sum.40, reduce_sum.41) +} + +region_6.9 { + reduce_sum.47 = f64[] parameter(0) + reduce_sum.48 = f64[] parameter(1) + ROOT reduce_sum.49 = f64[] add(reduce_sum.47, reduce_sum.48) +} + +region_7.10 { + reduce_sum.54 = f64[] parameter(0) + reduce_sum.55 = f64[] parameter(1) + ROOT reduce_sum.56 = f64[] add(reduce_sum.54, reduce_sum.55) +} + +ENTRY main.11 { + a_2_.1 = f64[2048]{0} parameter(2) + broadcast_in_dim.12 = f64[1,2048,1]{2,1,0} reshape(a_2_.1) + mul.42 = f64[1,2048,1]{2,1,0} broadcast(broadcast_in_dim.12), dimensions={0,1,2} + mul.43 = f64[2048]{0} reshape(mul.42) + mul.44 = f64[2048,2048,3]{2,1,0} broadcast(mul.43), dimensions={1} + mul.29 = f64[1,2048,1]{2,1,0} broadcast(broadcast_in_dim.12), dimensions={0,1,2} + mul.30 = f64[2048]{0} reshape(mul.29) + mul.31 = f64[2048,2048,3]{2,1,0} broadcast(mul.30), dimensions={1} + a_0_.1 = f64[2048,3]{1,0} parameter(0) + broadcast_in_dim.10 = f64[2048,1,3]{2,1,0} reshape(a_0_.1) + sub.7 = f64[2048,1,3]{2,1,0} broadcast(broadcast_in_dim.10), dimensions={0,1,2} + sub.8 = f64[2048,3]{1,0} reshape(sub.7) + sub.9 = f64[2048,2048,3]{2,1,0} broadcast(sub.8), dimensions={0,2} + a_1_.1 = f64[2048,3]{1,0} parameter(1) + broadcast_in_dim.11 = f64[1,2048,3]{2,1,0} reshape(a_1_.1) + sub.10 = f64[1,2048,3]{2,1,0} broadcast(broadcast_in_dim.11), dimensions={0,1,2} + sub.11 = f64[2048,3]{1,0} reshape(sub.10) + sub.12 = f64[2048,2048,3]{2,1,0} broadcast(sub.11), dimensions={1,2} + sub.13 = f64[2048,2048,3]{2,1,0} subtract(sub.9, sub.12) + mul.32 = f64[2048,2048,3]{2,1,0} multiply(mul.31, sub.13) + jvp_jit_norm__.3 = (f64[2048,2048,1]{2,1,0}, f64[2048,2048,1]{2,1,0}) call(sub.13), to_apply=norm.2 + jvp_jit_norm__.4 = f64[2048,2048,1]{2,1,0} get-tuple-element(jvp_jit_norm__.3), index=0 + integer_pow.3 = f64[2048,2048,1]{2,1,0} multiply(jvp_jit_norm__.4, jvp_jit_norm__.4) + integer_pow.4 = f64[2048,2048,1]{2,1,0} multiply(integer_pow.3, jvp_jit_norm__.4) + div.13 = f64[2048,2048,1]{2,1,0} broadcast(integer_pow.4), dimensions={0,1,2} + div.14 = f64[2048,2048]{1,0} reshape(div.13) + div.15 = f64[2048,2048,3]{2,1,0} broadcast(div.14), dimensions={0,1} + div.16 = f64[2048,2048,3]{2,1,0} divide(mul.32, div.15) + constant.13 = f64[] constant(0) + reduce_sum.58 = f64[2048,3]{1,0} reduce(div.16, constant.13), dimensions={1}, to_apply=region_1.3 + constant.10 = f64[] constant(2) + mul.27 = f64[2048,3]{1,0} broadcast(constant.10), dimensions={} + mul.35 = f64[2048,3]{1,0} multiply(reduce_sum.58, mul.27) + broadcast_in_dim.13 = f64[2048,2048,3]{2,1,0} broadcast(mul.35), dimensions={0,2} + div.18 = f64[2048,2048,1]{2,1,0} broadcast(integer_pow.4), dimensions={0,1,2} + div.19 = f64[2048,2048]{1,0} reshape(div.18) + div.20 = f64[2048,2048,3]{2,1,0} broadcast(div.19), dimensions={0,1} + div.21 = f64[2048,2048,3]{2,1,0} divide(broadcast_in_dim.13, div.20) + mul.45 = f64[2048,2048,3]{2,1,0} multiply(mul.44, div.21) + jvp_jit_norm__.5 = f64[2048,2048,1]{2,1,0} get-tuple-element(jvp_jit_norm__.3), index=1 + constant.11 = f64[] constant(1) + broadcast_in_dim.9 = f64[2048,2048,1]{2,1,0} broadcast(constant.11), dimensions={} + mul.34 = f64[2048,2048,1]{2,1,0} multiply(integer_pow.4, integer_pow.4) + div.17 = f64[2048,2048,1]{2,1,0} divide(broadcast_in_dim.9, mul.34) + mul.36 = f64[2048,2048,1]{2,1,0} broadcast(div.17), dimensions={0,1,2} + mul.37 = f64[2048,2048]{1,0} reshape(mul.36) + mul.38 = f64[2048,2048,3]{2,1,0} broadcast(mul.37), dimensions={0,1} + mul.39 = f64[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.13, mul.38) + mul.40 = f64[2048,2048,3]{2,1,0} multiply(mul.39, mul.32) + reduce_sum.59 = f64[2048,2048]{1,0} reduce(mul.40, constant.13), dimensions={2}, to_apply=region_2.4 + reshape.3 = f64[2048,2048,1]{2,1,0} reshape(reduce_sum.59) + neg.2 = f64[2048,2048,1]{2,1,0} negate(reshape.3) + integer_pow.5 = f64[2048,2048,1]{2,1,0} multiply(jvp_jit_norm__.4, jvp_jit_norm__.4) + constant.12 = f64[] constant(3) + mul.28 = f64[2048,2048,1]{2,1,0} broadcast(constant.12), dimensions={} + mul.33 = f64[2048,2048,1]{2,1,0} multiply(integer_pow.5, mul.28) + mul.41 = f64[2048,2048,1]{2,1,0} multiply(neg.2, mul.33) + transpose_jvp_jit_norm___.1 = f64[2048,2048,3]{2,1,0} call(sub.13, jvp_jit_norm__.5, mul.41), to_apply=norm_0.6 + add_any.3 = f64[2048,2048,3]{2,1,0} add(mul.45, transpose_jvp_jit_norm___.1) + reduce_sum.60 = f64[2048,3]{1,0} reduce(add_any.3, constant.13), dimensions={1}, to_apply=region_4.7 + reshape.4 = f64[2048,1,3]{2,1,0} reshape(reduce_sum.60) + reduce_sum.63 = f64[2048,3]{1,0} reduce(reshape.4, constant.13), dimensions={1}, to_apply=region_7.10 + neg.3 = f64[2048,2048,3]{2,1,0} negate(add_any.3) + reduce_sum.61 = f64[2048,3]{1,0} reduce(neg.3, constant.13), dimensions={0}, to_apply=region_5.8 + reshape.5 = f64[1,2048,3]{2,1,0} reshape(reduce_sum.61) + reduce_sum.62 = f64[2048,3]{1,0} reduce(reshape.5, constant.13), dimensions={0}, to_apply=region_6.9 + ROOT tuple.3 = (f64[2048,3]{1,0}, f64[2048,3]{1,0}) tuple(reduce_sum.63, reduce_sum.62) +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_min_distance_grad_2048_f64.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_min_distance_grad_2048_f64.hlo new file mode 100644 index 00000000000000..2504924e0e30a2 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_min_distance_grad_2048_f64.hlo @@ -0,0 +1,108 @@ +HloModule jit_loss, entry_computation_layout={(f64[2048,3]{1,0}, f64[2048,3]{1,0})->(f64[2048,3]{1,0}, f64[2048,3]{1,0})} + +region_0.1 { + reduce_sum.3 = f64[] parameter(0) + reduce_sum.4 = f64[] parameter(1) + ROOT reduce_sum.5 = f64[] add(reduce_sum.3, reduce_sum.4) +} + +norm.2 { + Arg_0.1 = f64[2048,2048,3]{2,1,0} parameter(0) + mul.1 = f64[2048,2048,3]{2,1,0} multiply(Arg_0.1, Arg_0.1) + constant.6 = f64[] constant(0) + reduce_sum.7 = f64[2048,2048]{1,0} reduce(mul.1, constant.6), dimensions={2}, to_apply=region_0.1 + sqrt.1 = f64[2048,2048]{1,0} sqrt(reduce_sum.7) + constant.5 = f64[] constant(0.5) + div.2 = f64[2048,2048]{1,0} broadcast(constant.5), dimensions={} + div.3 = f64[2048,2048]{1,0} divide(div.2, sqrt.1) + ROOT tuple.1 = (f64[2048,2048]{1,0}, f64[2048,2048]{1,0}) tuple(sqrt.1, div.3) +} + +region_1.3 { + reduce_min.3 = f64[] parameter(0) + reduce_min.4 = f64[] parameter(1) + ROOT reduce_min.5 = f64[] minimum(reduce_min.3, reduce_min.4) +} + +region_2.4 { + reduce_sum.11 = f64[] parameter(0) + reduce_sum.12 = f64[] parameter(1) + ROOT reduce_sum.13 = f64[] add(reduce_sum.11, reduce_sum.12) +} + +norm_0.5 { + Arg_0.3 = f64[2048,2048,3]{2,1,0} parameter(0) + Arg_2.1 = f64[2048,2048]{1,0} parameter(2) + Arg_1.1 = f64[2048,2048]{1,0} parameter(1) + mul.6 = f64[2048,2048]{1,0} multiply(Arg_2.1, Arg_1.1) + broadcast_in_dim.4 = f64[2048,2048,3]{2,1,0} broadcast(mul.6), dimensions={0,1} + mul.7 = f64[2048,2048,3]{2,1,0} multiply(Arg_0.3, broadcast_in_dim.4) + mul.8 = f64[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.4, Arg_0.3) + ROOT add_any.1 = f64[2048,2048,3]{2,1,0} add(mul.7, mul.8) +} + +region_3.6 { + reduce_sum.18 = f64[] parameter(0) + reduce_sum.19 = f64[] parameter(1) + ROOT reduce_sum.20 = f64[] add(reduce_sum.18, reduce_sum.19) +} + +region_4.7 { + reduce_sum.25 = f64[] parameter(0) + reduce_sum.26 = f64[] parameter(1) + ROOT reduce_sum.27 = f64[] add(reduce_sum.25, reduce_sum.26) +} + +region_5.8 { + reduce_sum.32 = f64[] parameter(0) + reduce_sum.33 = f64[] parameter(1) + ROOT reduce_sum.34 = f64[] add(reduce_sum.32, reduce_sum.33) +} + +region_6.9 { + reduce_sum.39 = f64[] parameter(0) + reduce_sum.40 = f64[] parameter(1) + ROOT reduce_sum.41 = f64[] add(reduce_sum.39, reduce_sum.40) +} + +ENTRY main.10 { + a_0_.1 = f64[2048,3]{1,0} parameter(0) + broadcast_in_dim.5 = f64[2048,1,3]{2,1,0} reshape(a_0_.1) + sub.7 = f64[2048,1,3]{2,1,0} broadcast(broadcast_in_dim.5), dimensions={0,1,2} + sub.8 = f64[2048,3]{1,0} reshape(sub.7) + sub.9 = f64[2048,2048,3]{2,1,0} broadcast(sub.8), dimensions={0,2} + a_1_.1 = f64[2048,3]{1,0} parameter(1) + broadcast_in_dim.6 = f64[1,2048,3]{2,1,0} reshape(a_1_.1) + sub.10 = f64[1,2048,3]{2,1,0} broadcast(broadcast_in_dim.6), dimensions={0,1,2} + sub.11 = f64[2048,3]{1,0} reshape(sub.10) + sub.12 = f64[2048,2048,3]{2,1,0} broadcast(sub.11), dimensions={1,2} + sub.13 = f64[2048,2048,3]{2,1,0} subtract(sub.9, sub.12) + jvp_jit_norm__.3 = (f64[2048,2048]{1,0}, f64[2048,2048]{1,0}) call(sub.13), to_apply=norm.2 + jvp_jit_norm__.5 = f64[2048,2048]{1,0} get-tuple-element(jvp_jit_norm__.3), index=1 + constant.7 = f64[] constant(1) + jvp_jit_norm__.4 = f64[2048,2048]{1,0} get-tuple-element(jvp_jit_norm__.3), index=0 + constant.9 = f64[] constant(inf) + reduce_min.7 = f64[] reduce(jvp_jit_norm__.4, constant.9), dimensions={0,1}, to_apply=region_1.3 + reshape.3 = f64[1,1]{1,0} reshape(reduce_min.7) + eq.4 = f64[1,1]{1,0} broadcast(reshape.3), dimensions={0,1} + eq.5 = f64[] reshape(eq.4) + eq.6 = f64[2048,2048]{1,0} broadcast(eq.5), dimensions={} + eq.7 = pred[2048,2048]{1,0} compare(jvp_jit_norm__.4, eq.6), direction=EQ + convert_element_type.1 = f64[2048,2048]{1,0} convert(eq.7) + constant.8 = f64[] constant(0) + reduce_sum.43 = f64[] reduce(convert_element_type.1, constant.8), dimensions={0,1}, to_apply=region_2.4 + div.5 = f64[] divide(constant.7, reduce_sum.43) + broadcast_in_dim.7 = f64[2048,2048]{1,0} broadcast(div.5), dimensions={} + mul.9 = f64[2048,2048]{1,0} multiply(broadcast_in_dim.7, convert_element_type.1) + transpose_jvp_jit_norm___.1 = f64[2048,2048,3]{2,1,0} call(sub.13, jvp_jit_norm__.5, mul.9), to_apply=norm_0.5 + reduce_sum.44 = f64[2048,3]{1,0} reduce(transpose_jvp_jit_norm___.1, constant.8), dimensions={1}, to_apply=region_3.6 + reshape.4 = f64[2048,1,3]{2,1,0} reshape(reduce_sum.44) + reduce_sum.47 = f64[2048,3]{1,0} reduce(reshape.4, constant.8), dimensions={1}, to_apply=region_6.9 + neg.1 = f64[2048,2048,3]{2,1,0} negate(transpose_jvp_jit_norm___.1) + reduce_sum.45 = f64[2048,3]{1,0} reduce(neg.1, constant.8), dimensions={0}, to_apply=region_4.7 + reshape.5 = f64[1,2048,3]{2,1,0} reshape(reduce_sum.45) + reduce_sum.46 = f64[2048,3]{1,0} reduce(reshape.5, constant.8), dimensions={0}, to_apply=region_5.8 + ROOT tuple.3 = (f64[2048,3]{1,0}, f64[2048,3]{1,0}) tuple(reduce_sum.47, reduce_sum.46) +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_4096_f32.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_4096_f32.hlo new file mode 100644 index 00000000000000..34842a9d95cd30 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_4096_f32.hlo @@ -0,0 +1,48 @@ +HloModule jit_rbf_gradient_field, entry_computation_layout={(f32[4096,3]{1,0}, f32[4096,3]{1,0}, f32[4096]{0})->f32[4096,3]{1,0}} + +region_0.1 { + reduce_sum.3 = f32[] parameter(0) + reduce_sum.4 = f32[] parameter(1) + ROOT reduce_sum.5 = f32[] add(reduce_sum.3, reduce_sum.4) +} + +region_1.2 { + reduce_sum.10 = f32[] parameter(0) + reduce_sum.11 = f32[] parameter(1) + ROOT reduce_sum.12 = f32[] add(reduce_sum.10, reduce_sum.11) +} + +ENTRY main.3 { + w.1 = f32[4096]{0} parameter(2) + broadcast_in_dim.7 = f32[1,4096,1]{2,1,0} reshape(w.1) + mul.14 = f32[1,4096,1]{2,1,0} broadcast(broadcast_in_dim.7), dimensions={0,1,2} + mul.15 = f32[4096,1]{1,0} reshape(mul.14) + mul.16 = f32[4096,4096,1]{2,1,0} broadcast(mul.15), dimensions={1,2} + pos.1 = f32[4096,3]{1,0} parameter(0) + broadcast_in_dim.4 = f32[4096,1,3]{2,1,0} reshape(pos.1) + sub.7 = f32[4096,1,3]{2,1,0} broadcast(broadcast_in_dim.4), dimensions={0,1,2} + sub.8 = f32[4096,3]{1,0} reshape(sub.7) + sub.9 = f32[4096,4096,3]{2,1,0} broadcast(sub.8), dimensions={0,2} + src.1 = f32[4096,3]{1,0} parameter(1) + broadcast_in_dim.5 = f32[1,4096,3]{2,1,0} reshape(src.1) + sub.10 = f32[1,4096,3]{2,1,0} broadcast(broadcast_in_dim.5), dimensions={0,1,2} + sub.11 = f32[4096,3]{1,0} reshape(sub.10) + sub.12 = f32[4096,4096,3]{2,1,0} broadcast(sub.11), dimensions={1,2} + sub.13 = f32[4096,4096,3]{2,1,0} subtract(sub.9, sub.12) + mul.12 = f32[4096,4096,3]{2,1,0} multiply(sub.13, sub.13) + constant.3 = f32[] constant(0) + reduce_sum.14 = f32[4096,4096]{1,0} reduce(mul.12, constant.3), dimensions={2}, to_apply=region_0.1 + broadcast_in_dim.6 = f32[4096,4096,1]{2,1,0} reshape(reduce_sum.14) + constant.2 = f32[] constant(-0.5) + mul.11 = f32[4096,4096,1]{2,1,0} broadcast(constant.2), dimensions={} + mul.13 = f32[4096,4096,1]{2,1,0} multiply(broadcast_in_dim.6, mul.11) + exp.1 = f32[4096,4096,1]{2,1,0} exponential(mul.13) + mul.17 = f32[4096,4096,1]{2,1,0} multiply(mul.16, exp.1) + mul.18 = f32[4096,4096,1]{2,1,0} broadcast(mul.17), dimensions={0,1,2} + mul.19 = f32[4096,4096]{1,0} reshape(mul.18) + mul.20 = f32[4096,4096,3]{2,1,0} broadcast(mul.19), dimensions={0,1} + mul.21 = f32[4096,4096,3]{2,1,0} multiply(mul.20, sub.13) + ROOT reduce_sum.15 = f32[4096,3]{1,0} reduce(mul.21, constant.3), dimensions={1}, to_apply=region_1.2 +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_grad_2048_f32.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_grad_2048_f32.hlo new file mode 100644 index 00000000000000..33990008aa1406 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_rbf_gradient_grad_2048_f32.hlo @@ -0,0 +1,115 @@ +HloModule jit_loss, entry_computation_layout={(f32[2048,3]{1,0}, f32[2048,3]{1,0}, f32[2048]{0})->(f32[2048,3]{1,0}, f32[2048,3]{1,0})} + +region_0.1 { + reduce_sum.3 = f32[] parameter(0) + reduce_sum.4 = f32[] parameter(1) + ROOT reduce_sum.5 = f32[] add(reduce_sum.3, reduce_sum.4) +} + +region_1.2 { + reduce_sum.10 = f32[] parameter(0) + reduce_sum.11 = f32[] parameter(1) + ROOT reduce_sum.12 = f32[] add(reduce_sum.10, reduce_sum.11) +} + +region_2.3 { + reduce_sum.17 = f32[] parameter(0) + reduce_sum.18 = f32[] parameter(1) + ROOT reduce_sum.19 = f32[] add(reduce_sum.17, reduce_sum.18) +} + +region_3.4 { + reduce_sum.24 = f32[] parameter(0) + reduce_sum.25 = f32[] parameter(1) + ROOT reduce_sum.26 = f32[] add(reduce_sum.24, reduce_sum.25) +} + +region_4.5 { + reduce_sum.31 = f32[] parameter(0) + reduce_sum.32 = f32[] parameter(1) + ROOT reduce_sum.33 = f32[] add(reduce_sum.31, reduce_sum.32) +} + +region_5.6 { + reduce_sum.38 = f32[] parameter(0) + reduce_sum.39 = f32[] parameter(1) + ROOT reduce_sum.40 = f32[] add(reduce_sum.38, reduce_sum.39) +} + +region_6.7 { + reduce_sum.45 = f32[] parameter(0) + reduce_sum.46 = f32[] parameter(1) + ROOT reduce_sum.47 = f32[] add(reduce_sum.45, reduce_sum.46) +} + +region_7.8 { + reduce_sum.52 = f32[] parameter(0) + reduce_sum.53 = f32[] parameter(1) + ROOT reduce_sum.54 = f32[] add(reduce_sum.52, reduce_sum.53) +} + +ENTRY main.9 { + a_2_.1 = f32[2048]{0} parameter(2) + broadcast_in_dim.9 = f32[1,2048,1]{2,1,0} reshape(a_2_.1) + mul.28 = f32[1,2048,1]{2,1,0} broadcast(broadcast_in_dim.9), dimensions={0,1,2} + mul.29 = f32[2048,1]{1,0} reshape(mul.28) + mul.30 = f32[2048,2048,1]{2,1,0} broadcast(mul.29), dimensions={1,2} + a_0_.1 = f32[2048,3]{1,0} parameter(0) + broadcast_in_dim.6 = f32[2048,1,3]{2,1,0} reshape(a_0_.1) + sub.7 = f32[2048,1,3]{2,1,0} broadcast(broadcast_in_dim.6), dimensions={0,1,2} + sub.8 = f32[2048,3]{1,0} reshape(sub.7) + sub.9 = f32[2048,2048,3]{2,1,0} broadcast(sub.8), dimensions={0,2} + a_1_.1 = f32[2048,3]{1,0} parameter(1) + broadcast_in_dim.7 = f32[1,2048,3]{2,1,0} reshape(a_1_.1) + sub.10 = f32[1,2048,3]{2,1,0} broadcast(broadcast_in_dim.7), dimensions={0,1,2} + sub.11 = f32[2048,3]{1,0} reshape(sub.10) + sub.12 = f32[2048,2048,3]{2,1,0} broadcast(sub.11), dimensions={1,2} + sub.13 = f32[2048,2048,3]{2,1,0} subtract(sub.9, sub.12) + mul.26 = f32[2048,2048,3]{2,1,0} multiply(sub.13, sub.13) + constant.5 = f32[] constant(0) + reduce_sum.56 = f32[2048,2048]{1,0} reduce(mul.26, constant.5), dimensions={2}, to_apply=region_0.1 + broadcast_in_dim.8 = f32[2048,2048,1]{2,1,0} reshape(reduce_sum.56) + constant.4 = f32[] constant(-0.5) + broadcast.1 = f32[2048,2048,1]{2,1,0} broadcast(constant.4), dimensions={} + mul.27 = f32[2048,2048,1]{2,1,0} multiply(broadcast_in_dim.8, broadcast.1) + exp.1 = f32[2048,2048,1]{2,1,0} exponential(mul.27) + mul.31 = f32[2048,2048,1]{2,1,0} multiply(mul.30, exp.1) + mul.37 = f32[2048,2048,1]{2,1,0} broadcast(mul.31), dimensions={0,1,2} + mul.38 = f32[2048,2048]{1,0} reshape(mul.37) + mul.39 = f32[2048,2048,3]{2,1,0} broadcast(mul.38), dimensions={0,1} + mul.32 = f32[2048,2048,1]{2,1,0} broadcast(mul.31), dimensions={0,1,2} + mul.33 = f32[2048,2048]{1,0} reshape(mul.32) + mul.34 = f32[2048,2048,3]{2,1,0} broadcast(mul.33), dimensions={0,1} + mul.35 = f32[2048,2048,3]{2,1,0} multiply(mul.34, sub.13) + reduce_sum.57 = f32[2048,3]{1,0} reduce(mul.35, constant.5), dimensions={1}, to_apply=region_1.2 + constant.3 = f32[] constant(2) + mul.25 = f32[2048,3]{1,0} broadcast(constant.3), dimensions={} + mul.36 = f32[2048,3]{1,0} multiply(reduce_sum.57, mul.25) + broadcast_in_dim.10 = f32[2048,2048,3]{2,1,0} broadcast(mul.36), dimensions={0,2} + mul.40 = f32[2048,2048,3]{2,1,0} multiply(mul.39, broadcast_in_dim.10) + mul.42 = f32[1,2048,1]{2,1,0} broadcast(broadcast_in_dim.9), dimensions={0,1,2} + mul.43 = f32[2048,1]{1,0} reshape(mul.42) + mul.44 = f32[2048,2048,1]{2,1,0} broadcast(mul.43), dimensions={1,2} + mul.41 = f32[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.10, sub.13) + reduce_sum.58 = f32[2048,2048]{1,0} reduce(mul.41, constant.5), dimensions={2}, to_apply=region_2.3 + reshape.3 = f32[2048,2048,1]{2,1,0} reshape(reduce_sum.58) + mul.45 = f32[2048,2048,1]{2,1,0} multiply(mul.44, reshape.3) + mul.46 = f32[2048,2048,1]{2,1,0} multiply(mul.45, exp.1) + mul.47 = f32[2048,2048,1]{2,1,0} multiply(mul.46, broadcast.1) + reduce_sum.59 = f32[2048,2048]{1,0} reduce(mul.47, constant.5), dimensions={2}, to_apply=region_3.4 + broadcast_in_dim.11 = f32[2048,2048,3]{2,1,0} broadcast(reduce_sum.59), dimensions={0,1} + mul.48 = f32[2048,2048,3]{2,1,0} multiply(sub.13, broadcast_in_dim.11) + add_any.2 = f32[2048,2048,3]{2,1,0} add(mul.40, mul.48) + mul.49 = f32[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.11, sub.13) + add_any.3 = f32[2048,2048,3]{2,1,0} add(add_any.2, mul.49) + reduce_sum.60 = f32[2048,3]{1,0} reduce(add_any.3, constant.5), dimensions={1}, to_apply=region_4.5 + reshape.4 = f32[2048,1,3]{2,1,0} reshape(reduce_sum.60) + reduce_sum.63 = f32[2048,3]{1,0} reduce(reshape.4, constant.5), dimensions={1}, to_apply=region_7.8 + neg.1 = f32[2048,2048,3]{2,1,0} negate(add_any.3) + reduce_sum.61 = f32[2048,3]{1,0} reduce(neg.1, constant.5), dimensions={0}, to_apply=region_5.6 + reshape.5 = f32[1,2048,3]{2,1,0} reshape(reduce_sum.61) + reduce_sum.62 = f32[2048,3]{1,0} reduce(reshape.5, constant.5), dimensions={0}, to_apply=region_6.7 + ROOT tuple.1 = (f32[2048,3]{1,0}, f32[2048,3]{1,0}) tuple(reduce_sum.63, reduce_sum.62) +} + + diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_sum_distances_grad_2048_f32.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_sum_distances_grad_2048_f32.hlo new file mode 100644 index 00000000000000..b383186fb7a62b --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/pairwise_sum_distances_grad_2048_f32.hlo @@ -0,0 +1,94 @@ +HloModule jit_loss, entry_computation_layout={(f32[2048,3]{1,0}, f32[2048,3]{1,0})->(f32[2048,3]{1,0}, f32[2048,3]{1,0})} + +region_0.1 { + reduce_sum.3 = f32[] parameter(0) + reduce_sum.4 = f32[] parameter(1) + ROOT reduce_sum.5 = f32[] add(reduce_sum.3, reduce_sum.4) +} + +norm.2 { + Arg_0.1 = f32[2048,2048,3]{2,1,0} parameter(0) + mul.2 = f32[2048,2048,3]{2,1,0} multiply(Arg_0.1, Arg_0.1) + constant.5 = f32[] constant(0) + reduce_sum.7 = f32[2048,2048]{1,0} reduce(mul.2, constant.5), dimensions={2}, to_apply=region_0.1 + sqrt.1 = f32[2048,2048]{1,0} sqrt(reduce_sum.7) + constant.4 = f32[] constant(0.5) + div.2 = f32[2048,2048]{1,0} broadcast(constant.4), dimensions={} + div.3 = f32[2048,2048]{1,0} divide(div.2, sqrt.1) + ROOT tuple.1 = (f32[2048,2048]{1,0}, f32[2048,2048]{1,0}) tuple(sqrt.1, div.3) +} + +region_1.3 { + reduce_sum.11 = f32[] parameter(0) + reduce_sum.12 = f32[] parameter(1) + ROOT reduce_sum.13 = f32[] add(reduce_sum.11, reduce_sum.12) +} + +norm_0.4 { + Arg_0.3 = f32[2048,2048,3]{2,1,0} parameter(0) + Arg_2.1 = f32[2048,2048]{1,0} parameter(2) + Arg_1.1 = f32[2048,2048]{1,0} parameter(1) + mul.7 = f32[2048,2048]{1,0} multiply(Arg_2.1, Arg_1.1) + broadcast_in_dim.4 = f32[2048,2048,3]{2,1,0} broadcast(mul.7), dimensions={0,1} + mul.8 = f32[2048,2048,3]{2,1,0} multiply(Arg_0.3, broadcast_in_dim.4) + mul.9 = f32[2048,2048,3]{2,1,0} multiply(broadcast_in_dim.4, Arg_0.3) + ROOT add_any.1 = f32[2048,2048,3]{2,1,0} add(mul.8, mul.9) +} + +region_2.5 { + reduce_sum.18 = f32[] parameter(0) + reduce_sum.19 = f32[] parameter(1) + ROOT reduce_sum.20 = f32[] add(reduce_sum.18, reduce_sum.19) +} + +region_3.6 { + reduce_sum.25 = f32[] parameter(0) + reduce_sum.26 = f32[] parameter(1) + ROOT reduce_sum.27 = f32[] add(reduce_sum.25, reduce_sum.26) +} + +region_4.7 { + reduce_sum.32 = f32[] parameter(0) + reduce_sum.33 = f32[] parameter(1) + ROOT reduce_sum.34 = f32[] add(reduce_sum.32, reduce_sum.33) +} + +region_5.8 { + reduce_sum.39 = f32[] parameter(0) + reduce_sum.40 = f32[] parameter(1) + ROOT reduce_sum.41 = f32[] add(reduce_sum.39, reduce_sum.40) +} + +ENTRY main.9 { + a_0_.1 = f32[2048,3]{1,0} parameter(0) + broadcast_in_dim.5 = f32[2048,1,3]{2,1,0} reshape(a_0_.1) + sub.7 = f32[2048,1,3]{2,1,0} broadcast(broadcast_in_dim.5), dimensions={0,1,2} + sub.8 = f32[2048,3]{1,0} reshape(sub.7) + sub.9 = f32[2048,2048,3]{2,1,0} broadcast(sub.8), dimensions={0,2} + a_1_.1 = f32[2048,3]{1,0} parameter(1) + broadcast_in_dim.6 = f32[1,2048,3]{2,1,0} reshape(a_1_.1) + sub.10 = f32[1,2048,3]{2,1,0} broadcast(broadcast_in_dim.6), dimensions={0,1,2} + sub.11 = f32[2048,3]{1,0} reshape(sub.10) + sub.12 = f32[2048,2048,3]{2,1,0} broadcast(sub.11), dimensions={1,2} + sub.13 = f32[2048,2048,3]{2,1,0} subtract(sub.9, sub.12) + jvp_jit_norm__.3 = (f32[2048,2048]{1,0}, f32[2048,2048]{1,0}) call(sub.13), to_apply=norm.2 + jvp_jit_norm__.5 = f32[2048,2048]{1,0} get-tuple-element(jvp_jit_norm__.3), index=1 + jvp_jit_norm__.4 = f32[2048,2048]{1,0} get-tuple-element(jvp_jit_norm__.3), index=0 + constant.7 = f32[] constant(0) + reduce_sum.43 = f32[2048]{0} reduce(jvp_jit_norm__.4, constant.7), dimensions={1}, to_apply=region_1.3 + constant.6 = f32[] constant(2) + mul.10 = f32[2048]{0} broadcast(constant.6), dimensions={} + mul.11 = f32[2048]{0} multiply(reduce_sum.43, mul.10) + broadcast_in_dim.7 = f32[2048,2048]{1,0} broadcast(mul.11), dimensions={0} + transpose_jvp_jit_norm___.1 = f32[2048,2048,3]{2,1,0} call(sub.13, jvp_jit_norm__.5, broadcast_in_dim.7), to_apply=norm_0.4 + reduce_sum.44 = f32[2048,3]{1,0} reduce(transpose_jvp_jit_norm___.1, constant.7), dimensions={1}, to_apply=region_2.5 + reshape.2 = f32[2048,1,3]{2,1,0} reshape(reduce_sum.44) + reduce_sum.47 = f32[2048,3]{1,0} reduce(reshape.2, constant.7), dimensions={1}, to_apply=region_5.8 + neg.1 = f32[2048,2048,3]{2,1,0} negate(transpose_jvp_jit_norm___.1) + reduce_sum.45 = f32[2048,3]{1,0} reduce(neg.1, constant.7), dimensions={0}, to_apply=region_3.6 + reshape.3 = f32[1,2048,3]{2,1,0} reshape(reduce_sum.45) + reduce_sum.46 = f32[2048,3]{1,0} reduce(reshape.3, constant.7), dimensions={0}, to_apply=region_4.7 + ROOT tuple.3 = (f32[2048,3]{1,0}, f32[2048,3]{1,0}) tuple(reduce_sum.47, reduce_sum.46) +} + + diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD index 65c152c3e6ea0a..47882a9e7826ae 100644 --- a/third_party/xla/xla/backends/cpu/codegen/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/BUILD @@ -105,6 +105,15 @@ cc_library( ], ) +cc_library( + name = "object_buffer_identifier", + hdrs = ["object_buffer_identifier.h"], + deps = [ + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + ], +) + cc_library( name = "ir_compiler", srcs = ["ir_compiler.cc"], @@ -120,6 +129,7 @@ cc_library( ]), deps = [ ":kernel_api_ir_builder", + ":object_buffer_identifier", ":polynomial_approximations", "//xla:util", "//xla:xla_proto_cc", @@ -167,6 +177,7 @@ xla_cc_test( deps = [ ":ir_compiler", ":kernel_api_ir_builder", + ":object_buffer_identifier", "//xla:debug_options_flags", "//xla:util", "//xla/backends/cpu:target_machine_options", @@ -609,11 +620,13 @@ cc_library( deps = [ ":contiguous_section_memory_manager", ":jit_memory_mapper", + ":object_buffer_identifier", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", "@llvm-project//llvm:ExecutionEngine", "@llvm-project//llvm:OrcJIT", "@llvm-project//llvm:Support", @@ -761,6 +774,7 @@ cc_library( "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/numeric:bits", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter_test.py b/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter_test.py index 60740906098e6f..5bd583b231e45e 100644 --- a/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter_test.py +++ b/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter_test.py @@ -48,7 +48,6 @@ def create_input( emitter_types = [ - testlib_cpu.ElementalKernelEmitter, testlib_cpu.DotKernelEmitter, ] diff --git a/third_party/xla/xla/backends/cpu/codegen/elemental/BUILD b/third_party/xla/xla/backends/cpu/codegen/elemental/BUILD index 3c8b10b896452a..8243e7f917f352 100644 --- a/third_party/xla/xla/backends/cpu/codegen/elemental/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/elemental/BUILD @@ -14,7 +14,6 @@ # ============================================================================== load("@xla//third_party/rules_python/python:py_test.bzl", "py_test") -load("//xla:xla.default.bzl", "xla_cc_test") load("//xla/tsl/platform:rules_cc.bzl", "cc_library") package( @@ -35,7 +34,6 @@ cc_library( srcs = ["concatenate_kernel_emitter.cc"], hdrs = ["concatenate_kernel_emitter.h"], deps = [ - ":elemental_kernel_emitter", "//xla:shape_util", "//xla:util", "//xla/backends/cpu/codegen:kernel_api_ir_builder", @@ -52,7 +50,6 @@ cc_library( "//xla/service/llvm_ir:ir_array", "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", @@ -76,67 +73,3 @@ py_test( "@absl_py//absl/testing:parameterized", ], ) - -cc_library( - name = "elemental_kernel_emitter", - srcs = ["elemental_kernel_emitter.cc"], - hdrs = ["elemental_kernel_emitter.h"], - deps = [ - "//xla:shape_util", - "//xla:util", - "//xla/backends/cpu/codegen:kernel_api_ir_builder", - "//xla/backends/cpu/codegen:target_machine_features", - "//xla/codegen:kernel_definition", - "//xla/codegen:kernel_emitter", - "//xla/codegen:kernel_spec", - "//xla/codegen:llvm_kernel_source", - "//xla/hlo/ir:hlo", - "//xla/runtime:work_group", - "//xla/service:buffer_assignment", - "//xla/service:elemental_ir_emitter", - "//xla/service:hlo_module_config", - "//xla/service/cpu:backend_config_proto_cc", - "//xla/service/cpu:elemental_ir_emitter", - "//xla/service/cpu:ir_emitter", - "//xla/service/cpu:parallel_loop_emitter", - "//xla/service/llvm_ir:ir_array", - "//xla/service/llvm_ir:llvm_util", - "//xla/service/llvm_ir:loop_emitter", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/log", - "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - "@llvm-project//llvm:ir_headers", - ], -) - -xla_cc_test( - name = "elemental_kernel_emitter_test", - srcs = ["elemental_kernel_emitter_test.cc"], - deps = [ - ":elemental_kernel_emitter", - "//xla:xla_data_proto_cc", - "//xla/backends/cpu/codegen:target_machine_test_base", - "//xla/codegen:kernel_definition", - "//xla/codegen:llvm_kernel_source", - "//xla/hlo/analysis:alias_info", - "//xla/hlo/analysis:hlo_ordering", - "//xla/hlo/ir:hlo", - "//xla/hlo/parser:hlo_parser", - "//xla/hlo/testlib:filecheck", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "//xla/service:buffer_assignment", - "//xla/service:buffer_value", - "//xla/service:logical_buffer", - "//xla/service/cpu:cpu_executable", - "//xla/tsl/platform:statusor", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - "@llvm-project//llvm:ir_headers", - ], -) diff --git a/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.cc b/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.cc index a877fa3bb3670d..955709471c4cc9 100644 --- a/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.cc +++ b/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.cc @@ -21,7 +21,6 @@ limitations under the License. #include #include "absl/algorithm/container.h" -#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -29,7 +28,6 @@ limitations under the License. #include "llvm/IR/Analysis.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/LLVMContext.h" -#include "xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h" #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" #include "xla/backends/cpu/codegen/target_machine_features.h" #include "xla/codegen/kernel_definition.h" @@ -70,10 +68,9 @@ ConcatenateKernelEmitter::ConcatenateKernelEmitter( absl::StatusOr ConcatenateKernelEmitter::EmitKernelDefinition() { if (absl::Status status = CanDoFastConcatenate(*instr_); !status.ok()) { - VLOG(1) << "Could not emit fast concatenate for " << instr_->ToString() - << ": " << status.message(); - return ElementalKernelEmitter(instr_, buffer_assignment_, target_machine_) - .EmitKernelDefinition(); + return Internal( + "Concatenate is not supported by ConcatenateKernelEmitter: %s", + status.message()); } auto ctx = std::make_unique(); @@ -83,8 +80,9 @@ ConcatenateKernelEmitter::EmitKernelDefinition() { return Internal("HloModule is null"); } - const auto& backend_config = instr_->backend_config(); - const auto& partitions = backend_config->outer_dimension_partitions(); + const auto backend_config = + instr_->backend_config().value_or(BackendConfig()); + const auto& partitions = backend_config.outer_dimension_partitions(); auto total_workgroups = absl::c_accumulate(partitions, 1, std::multiplies()); diff --git a/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter_test.py b/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter_test.py index 937540e36b9ee5..dcfbbb2b27dbd5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter_test.py +++ b/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter_test.py @@ -54,7 +54,8 @@ def test_concatenate(self, cycle_layout, dtype, concat_dimension): (np.random.rand(*shape) * 10).astype(dtype) for _ in range(num_inputs) ] if cycle_layout: - # Give the inputs different layouts to test the slow path. + # Give the inputs different layouts to test that mismatched layouts are + # rejected. default_layout = [0, 1, 2] input_literals = [ create_literal(input_array, np.roll(default_layout, idx)) @@ -88,6 +89,14 @@ def test_concatenate(self, cycle_layout, dtype, concat_dimension): jit_compiler.get_target_machine(), ) + if cycle_layout: + with self.assertRaisesRegex( + RuntimeError, + "Concatenate is not supported by ConcatenateKernelEmitter", + ): + emitter.emit_kernel_definition() + return + kernel_definition = emitter.emit_kernel_definition() self.assertIsNotNone(kernel_definition) diff --git a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.cc b/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.cc deleted file mode 100644 index e021773869e8dd..00000000000000 --- a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.cc +++ /dev/null @@ -1,340 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h" - -#include -#include -#include -#include -#include -#include - -#include "absl/container/flat_hash_map.h" -#include "absl/log/log.h" -#include "absl/status/status_macros.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "llvm/IR/Constants.h" -#include "llvm/IR/DerivedTypes.h" -#include "llvm/IR/IRBuilder.h" -#include "llvm/IR/Intrinsics.h" -#include "llvm/IR/LLVMContext.h" -#include "llvm/IR/Module.h" -#include "llvm/IR/Value.h" -#include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" -#include "xla/backends/cpu/codegen/target_machine_features.h" -#include "xla/codegen/kernel_definition.h" -#include "xla/codegen/kernel_spec.h" -#include "xla/codegen/llvm_kernel_source.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/runtime/work_group.h" -#include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/backend_config.pb.h" -#include "xla/service/cpu/elemental_ir_emitter.h" -#include "xla/service/cpu/ir_emitter.h" -#include "xla/service/cpu/parallel_loop_emitter.h" -#include "xla/service/elemental_ir_emitter.h" -#include "xla/service/hlo_module_config.h" -#include "xla/service/llvm_ir/ir_array.h" -#include "xla/service/llvm_ir/llvm_util.h" -#include "xla/service/llvm_ir/loop_emitter.h" -#include "xla/shape.h" -#include "xla/shape_partition.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" -#include "xla/util.h" - -namespace xla::cpu { - -namespace { - -struct ParallelConfig { - std::vector outer_dimension_partitions; -}; - -// Parallel partition bounds for parallelized outer dimensions: -// vector<[i64 lower_bound, i64 upper_bound]> -using ParallelPartitionBounds = - std::vector>; - -std::optional GetParallelConfig(const HloInstruction* instr) { - // Check if the instruction is marked for parallel execution. - auto backend_config = instr->backend_config(); - if (!backend_config.ok() || - backend_config->outer_dimension_partitions().empty()) { - return std::nullopt; - } - - ParallelConfig config; - config.outer_dimension_partitions.assign( - backend_config->outer_dimension_partitions().begin(), - backend_config->outer_dimension_partitions().end()); - - return config; -} - -ParallelPartitionBounds EmitParallelPartitionBounds( - llvm::IRBuilderBase& b, - const KernelApiIrBuilder::KernelPrototype& kernel_prototype, - const ParallelConfig& parallel_config, const Shape& shape, - absl::string_view name) { - ShapePartitionIterator it(shape, parallel_config.outer_dimension_partitions); - - size_t num_parallel_dimensions = - parallel_config.outer_dimension_partitions.size(); - - // Create a constant array of all partition bounds. We will be indexing into - // this array using workgroup id passed in a call frame. - // - // Type: [#partitions x [#outer_dimensions x [lower_bound, upper_bound]]] - // - llvm::ArrayType* dim_bounds_ty = llvm::ArrayType::get(b.getInt64Ty(), 2); - llvm::ArrayType* partition_bounds_ty = - llvm::ArrayType::get(dim_bounds_ty, num_parallel_dimensions); - llvm::ArrayType* parallel_bounds_ty = - llvm::ArrayType::get(partition_bounds_ty, it.GetTotalPartitionCount()); - - // Build a nested array of partition bounds from shape partition iterator. - std::vector partition_bounds; - for (int64_t i = 0; i < it.GetTotalPartitionCount(); ++i) { - std::vector dim_counts; - for (auto [lower, size] : it.GetPartition(i)) { - dim_counts.push_back(llvm::ConstantArray::get( - dim_bounds_ty, {b.getInt64(lower), b.getInt64(lower + size)})); - } - partition_bounds.push_back( - llvm::ConstantArray::get(partition_bounds_ty, dim_counts)); - } - - llvm::Constant* parallel_bounds = - llvm::ConstantArray::get(parallel_bounds_ty, partition_bounds); - - llvm::Module* module = b.GetInsertBlock()->getParent()->getParent(); - llvm::GlobalVariable* parallel_bounds_global = new llvm::GlobalVariable( - /*M=*/*module, - /*Ty=*/parallel_bounds_ty, - /*isConstant=*/true, - /*Linkage=*/llvm::GlobalValue::PrivateLinkage, - /*Initializer=*/parallel_bounds, - /*Name=*/absl::StrCat(name, "_parallel_bounds")); - - // Construct IR to load bounds for all parallel dimensions. - ParallelPartitionBounds bounds; - for (size_t i = 0; i < num_parallel_dimensions; ++i) { - llvm::Value* partition = kernel_prototype.workgroup_id.x; - llvm::Value* parallel_dim = b.getInt32(i); - - llvm::Value* lower_gep = b.CreateInBoundsGEP( - parallel_bounds_ty, parallel_bounds_global, - {b.getInt32(0), partition, parallel_dim, b.getInt32(0)}, - absl::StrCat("lo_dim_", i, "_gep")); - - llvm::Value* upper_gep = b.CreateInBoundsGEP( - parallel_bounds_ty, parallel_bounds_global, - {b.getInt32(0), partition, parallel_dim, b.getInt32(1)}, - absl::StrCat("up_dim_", i, "_gep")); - - bounds.emplace_back( - b.CreateLoad(b.getInt64Ty(), lower_gep, absl::StrCat("lo_dim_", i)), - b.CreateLoad(b.getInt64Ty(), upper_gep, absl::StrCat("up_dim_", i))); - } - - return bounds; -} - -} // namespace - -ElementalKernelEmitter::ElementalKernelEmitter( - const HloInstruction* instr, const BufferAssignment* buffer_assignment, - const TargetMachineFeatures* target_machine) - : instr_(instr), - buffer_assignment_(buffer_assignment), - target_machine_(target_machine) {} - -absl::StatusOr -ElementalKernelEmitter::EmitKernelDefinition() { - VLOG(2) << "Emit elemental host kernel: " << instr_->name(); - - auto ctx = std::make_unique(); - - const HloModule* hlo_module = instr_->GetModule(); - if (hlo_module == nullptr) { - return Internal("HloModule is null"); - } - - KernelApiIrBuilder kernel_api_ir_builder( - *ctx, - KernelApiIrBuilder::Options::FromHloModuleConfig(hlo_module->config())); - - std::unique_ptr llvm_module = KernelApiIrBuilder::CreateModule( - absl::StrCat(instr_->name(), "_elemental_kernel_module"), *ctx, - target_machine_); - - ABSL_ASSIGN_OR_RETURN( - KernelApiIrBuilder::KernelPrototype kernel_prototype, - kernel_api_ir_builder.EmitKernelPrototype( - *llvm_module, instr_, buffer_assignment_, name(), "_kernel")); - - llvm::IRBuilder<> ir_builder(*ctx); - ir_builder.SetInsertPoint( - kernel_prototype.function->getEntryBlock().getTerminator()); - - // Set fast-math flags on the builder. - ir_builder.setFastMathFlags( - llvm_ir::GetCpuFastMathFlags(hlo_module->config())); - - ABSL_ASSIGN_OR_RETURN( - CpuElementalIrEmitter::ThreadLocalCallCallback thread_local_call_fn, - ThreadLocalCallbackFactory(ir_builder, *llvm_module)); - - CpuElementalIrEmitter::HloToElementGeneratorMap operand_to_generator; - for (int64_t i = 0; i < instr_->operand_count(); ++i) { - const HloInstruction* operand = instr_->operand(i); - operand_to_generator[operand] = [&, i](const llvm_ir::IrArray::Index& idx) { - return kernel_prototype.arguments[i].EmitReadArrayElement(idx, - &ir_builder); - }; - } - - bool enable_fast_min_max = - hlo_module - ? hlo_module->config().debug_options().xla_cpu_enable_fast_min_max() - : true; - CpuElementalIrEmitter elemental_ir_emitter(llvm_module.get(), &ir_builder, - std::move(thread_local_call_fn), - true, enable_fast_min_max); - - llvm_ir::ElementGenerator element_generator = - elemental_ir_emitter.MakeElementGenerator(instr_, operand_to_generator); - - ABSL_ASSIGN_OR_RETURN(NumWorkGroups num_workgroups, - EmitElementalLoops(ir_builder, instr_, kernel_prototype, - element_generator)); - - LlvmKernelSource source(std::move(ctx), std::move(llvm_module)); - - KernelSpec::Buffers arguments, results; - for (const auto& buffer : kernel_prototype.argument_buffers) { - arguments.push_back({buffer.slice, buffer.shape}); - } - for (const auto& buffer : kernel_prototype.result_buffers) { - results.push_back({buffer.slice, buffer.shape}); - } - - KernelSpec spec(kernel_prototype.function->getName(), num_workgroups, - std::move(arguments), std::move(results), - std::move(kernel_prototype.invariant_arguments)); - - return KernelDefinition(std::move(spec), std::move(source)); -} - -absl::StatusOr ElementalKernelEmitter::EmitElementalLoops( - llvm::IRBuilderBase& b, const HloInstruction* instr, - const KernelApiIrBuilder::KernelPrototype& kernel_prototype, - const llvm_ir::ElementGenerator& element_generator) { - // We can emit loops for instruction with multiple results only if it is a - // fusion, reduce or reduce window. - bool multiple_results = kernel_prototype.results.size() > 1; - bool support_multiple_results = instr->opcode() == HloOpcode::kFusion || - instr->opcode() == HloOpcode::kReduce || - instr->opcode() == HloOpcode::kReduceWindow; - - auto parallel_config = GetParallelConfig(instr); - bool has_parallel_config = parallel_config.has_value(); - - if (multiple_results && !support_multiple_results) { - return Internal( - "Multi-output host kernels are not supported for %s instruction", - HloOpcodeString(instr->opcode())); - } - - // TODO(ezhulenev): Support multiple results for parallel loops. - if (multiple_results) { - ABSL_RETURN_IF_ERROR( - llvm_ir::LoopEmitter(element_generator, kernel_prototype.results, &b) - .EmitLoop(llvm_ir::IrName(instr))); - return NumWorkGroups(); - } - - const llvm_ir::IrArray& result = kernel_prototype.results.front(); - - // Emit a loop for a single parallel partition with dynamic bounds computed - // from workgroup id. - if (has_parallel_config) { - ParallelPartitionBounds parallel_bounds = EmitParallelPartitionBounds( - b, kernel_prototype, *parallel_config, instr->shape(), instr->name()); - ABSL_RETURN_IF_ERROR( - ParallelLoopEmitter(element_generator, result, ¶llel_bounds, &b) - .EmitLoop(llvm_ir::IrName(instr))); - return NumWorkGroups{ - static_cast(ShapePartitionAssigner::GetTotalPartitionCount( - parallel_config->outer_dimension_partitions))}; - } - - // Emit a whole loop for the instruction. - ABSL_RETURN_IF_ERROR(llvm_ir::LoopEmitter(element_generator, result, &b) - .EmitLoop(llvm_ir::IrName(instr))); - return NumWorkGroups(); -} - -absl::StatusOr -ElementalKernelEmitter::ThreadLocalCallbackFactory(llvm::IRBuilderBase& builder, - llvm::Module& module) const { - const HloModule* hlo_module = instr_->GetModule(); - if (hlo_module == nullptr) { - return nullptr; - } - - auto ir_emitter = std::make_unique( - nullptr, *hlo_module, *buffer_assignment_, &module, - /*instruction_to_profile_idx=*/ - absl::flat_hash_map{}, - /*computation_to_profile_idx=*/ - absl::flat_hash_map{}, - ComputationsTransitivelyContainCustomCall(instr_), target_machine_, - /*emit_code_for_msan=*/false); - IrEmitter::IRBuilderGuard builder_guard = ir_emitter->WithBuilder(builder); - - ABSL_RETURN_IF_ERROR(ir_emitter->EmitSmallConstantGlobals()); - - if (instr_->has_to_apply()) { - HloComputation* nested_computation = instr_->to_apply(); - bool is_reducer = instr_->opcode() == HloOpcode::kReduce || - instr_->opcode() == HloOpcode::kReduceWindow; - ABSL_RETURN_IF_ERROR( - ir_emitter - ->EmitNestedComputation(*nested_computation, - llvm_ir::IrName(nested_computation->name()), - is_reducer) - .status()); - } - - return [ir_emitter = std::move(ir_emitter), &builder]( - const HloComputation& callee, - absl::Span parameters, absl::string_view name, - bool is_reducer) { - IrEmitter::IRBuilderGuard builder_guard = ir_emitter->WithBuilder(builder); - return ir_emitter->EmitThreadLocalCall(callee, parameters, name, is_reducer, - /*in_compute_function=*/false); - }; -} - -} // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h b/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h deleted file mode 100644 index 59549f300719dc..00000000000000 --- a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h +++ /dev/null @@ -1,68 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_BACKENDS_CPU_CODEGEN_ELEMENTAL_ELEMENTAL_KERNEL_EMITTER_H_ -#define XLA_BACKENDS_CPU_CODEGEN_ELEMENTAL_ELEMENTAL_KERNEL_EMITTER_H_ - -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "llvm/IR/IRBuilder.h" -#include "llvm/IR/Module.h" -#include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" -#include "xla/backends/cpu/codegen/target_machine_features.h" -#include "xla/codegen/kernel_emitter.h" -#include "xla/codegen/llvm_kernel_source.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/runtime/work_group.h" -#include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/elemental_ir_emitter.h" -#include "xla/service/llvm_ir/loop_emitter.h" - -namespace xla::cpu { - -class ElementalKernelEmitter final : public KernelEmitter { - public: - ElementalKernelEmitter(const HloInstruction* instr, - const BufferAssignment* buffer_assignment, - const TargetMachineFeatures* target_machine); - - absl::string_view name() const final { return "elemental_kernel_emitter"; } - absl::StatusOr EmitKernelDefinition() override; - - private: - // Emits LLVM IR using elemental loop emitter and the given element generator. - // If the instruction is parallelized, it will emit a parallel loop partition - // and return the requested number of execution threads. - absl::StatusOr EmitElementalLoops( - llvm::IRBuilderBase& b, const HloInstruction* instr, - const KernelApiIrBuilder::KernelPrototype& kernel_prototype, - const llvm_ir::ElementGenerator& element_generator); - - // Create a thread local call callback, can be empty if no IrEmitter is - // registered. - absl::StatusOr - ThreadLocalCallbackFactory(llvm::IRBuilderBase& builder, - llvm::Module& module) const; - - private: - const HloInstruction* instr_; - - const BufferAssignment* buffer_assignment_ = nullptr; - const TargetMachineFeatures* target_machine_ = nullptr; -}; - -} // namespace xla::cpu - -#endif // XLA_BACKENDS_CPU_CODEGEN_ELEMENTAL_ELEMENTAL_KERNEL_EMITTER_H_ diff --git a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter_test.cc b/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter_test.cc deleted file mode 100644 index b7fb97f4956ed3..00000000000000 --- a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter_test.cc +++ /dev/null @@ -1,179 +0,0 @@ -/* Copyright 2024 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h" - -#include -#include - -#include -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "llvm/IR/LLVMContext.h" -#include "llvm/IR/Type.h" -#include "xla/backends/cpu/codegen/target_machine_test_base.h" -#include "xla/codegen/kernel_definition.h" -#include "xla/codegen/llvm_kernel_source.h" -#include "xla/hlo/analysis/alias_info.h" -#include "xla/hlo/analysis/hlo_ordering.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/parser/hlo_parser.h" -#include "xla/hlo/testlib/filecheck.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/service/buffer_assignment.h" -#include "xla/service/buffer_value.h" -#include "xla/service/cpu/cpu_executable.h" -#include "xla/service/logical_buffer.h" -#include "xla/tsl/platform/statusor.h" -#include "xla/xla_data.pb.h" - -namespace xla::cpu { - -class ElementalKernelEmitterTest : public TargetMachineTestBase { - protected: - void SetUp() override { - // register X86/ARM targets - TargetMachineTestBase::SetUp(); - target_machine_features_ = - CreateTargetMachineFeatures("x86_64-unknown-linux-gnu", "", ""); - } - - public: - absl::StatusOr> EmitKernelDefinition( - const HloInstruction* instr, const BufferAssignment* buffer_assignment) { - ElementalKernelEmitter emitter(instr, buffer_assignment, - target_machine_features_.get()); - - return emitter.EmitKernelDefinition(); - } - - absl::StatusOr> RunBufferAssignment( - const HloModule& hlo) { - return BufferAssigner::Run( - &hlo, std::make_unique(&hlo), - [](const BufferValue& buffer) { - return CpuExecutable::ShapeSizeBytes(buffer.shape()); - }, - &alias_info_, [](LogicalBuffer::Color) { return /*alignment=*/1; }, - BufferAssigner::Options{}); - } - - private: - std::unique_ptr target_machine_features_; - AliasInfo alias_info_; -}; - -namespace { - -TEST_F(ElementalKernelEmitterTest, EmitElementalKernel) { - const char* hlo_text = R"( - HloModule m - ENTRY main { - p0 = f32[2,2] parameter(0) - ROOT convert = s32[2,2] convert(p0) - })"; - - TF_ASSERT_OK_AND_ASSIGN(auto hlo, ParseAndReturnUnverifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN(auto buffer_assignment, RunBufferAssignment(*hlo)); - TF_ASSERT_OK_AND_ASSIGN( - KernelDefinition kernel_definition, - EmitKernelDefinition(hlo->entry_computation()->root_instruction(), - buffer_assignment.get())); - - ASSERT_TRUE(*RunFileCheck(kernel_definition.source().ToString(), R"( - CHECK: define ptr @convert_kernel(ptr noundef nonnull %0) #0 { - CHECK: fptosi float {{.*}} to i32 - CHECK: } - )")); -} - -TEST_F(ElementalKernelEmitterTest, EmitParallelKernel) { - llvm::LLVMContext context; - auto module = std::make_unique("test", context); - - const char* hlo_text = R"( - HloModule m - ENTRY main { - p0 = f32[1,2,1,16384,256] parameter(0) - ROOT convert = s32[1,2,1,16384,256] convert(p0), - backend_config={"outer_dimension_partitions":["1","2","1","4"]} - })"; - - TF_ASSERT_OK_AND_ASSIGN(auto hlo, ParseAndReturnUnverifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN(auto buffer_assignment, RunBufferAssignment(*hlo)); - TF_ASSERT_OK_AND_ASSIGN( - KernelDefinition kernel_definition, - EmitKernelDefinition(hlo->entry_computation()->root_instruction(), - buffer_assignment.get())); - - ASSERT_TRUE(*RunFileCheck(kernel_definition.source().ToString(), R"( - CHECK: @convert_parallel_bounds = private constant [8 x [4 x [2 x i64]]] - - CHECK: define ptr @convert_kernel(ptr noundef nonnull %0) #0 { - CHECK: %[[X:.*]] = load i64, ptr %workgroup_id_x_gep, align 8 - CHECK: %lo_dim_0_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 0, i32 0 - CHECK: %up_dim_0_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 0, i32 1 - CHECK: %lo_dim_1_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 1, i32 0 - CHECK: %up_dim_1_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 1, i32 1 - CHECK: %lo_dim_2_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 2, i32 0 - CHECK: %up_dim_2_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 2, i32 1 - CHECK: %lo_dim_3_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 3, i32 0 - CHECK: %up_dim_3_gep = getelementptr{{.*}} i32 0, i64 %[[X]], i32 3, i32 1 - CHECK: fptosi float {{.*}} to i32 - CHECK: } - )")); -} - -TEST_F(ElementalKernelEmitterTest, EmitFastIntrinsic) { - const char* hlo_text = R"( - HloModule m - ENTRY main { - p0 = f32[1024] parameter(0) - ROOT log = f32[1024] log(p0) - })"; - - TF_ASSERT_OK_AND_ASSIGN(auto hlo, ParseAndReturnUnverifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN(auto buffer_assignment, RunBufferAssignment(*hlo)); - - // Configure math options. - hlo->mutable_config().mutable_debug_options().set_xla_cpu_enable_fast_math( - true); - hlo->mutable_config() - .mutable_debug_options() - .set_xla_cpu_fast_math_honor_nans(false); - hlo->mutable_config() - .mutable_debug_options() - .set_xla_cpu_fast_math_honor_infs(false); - hlo->mutable_config() - .mutable_debug_options() - .set_xla_cpu_fast_math_honor_division(false); - hlo->mutable_config() - .mutable_debug_options() - .set_xla_cpu_fast_math_honor_functions(false); - - TF_ASSERT_OK_AND_ASSIGN( - KernelDefinition kernel_definition, - EmitKernelDefinition(hlo->entry_computation()->root_instruction(), - buffer_assignment.get())); - - ASSERT_TRUE(*RunFileCheck(kernel_definition.source().ToString(), R"( - CHECK: define ptr @log_kernel(ptr noundef nonnull %0) #0 { - CHECK: call fast float @llvm.log.f32(float {{.*}}) - CHECK: } - )")); -} - -} // namespace -} // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD b/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD index 7002e80d7b8cc4..0999c8517e5028 100644 --- a/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD @@ -119,6 +119,7 @@ xla_cc_test( deps = [ ":cpu_fusion_emitters", "//xla/backends/cpu/codegen:fusion_compiler", + "//xla/backends/cpu/codegen:fusion_emitter", "//xla/codegen:kernel_definition", "//xla/codegen:llvm_kernel_source", "//xla/codegen:mlir_kernel_source", @@ -127,6 +128,7 @@ xla_cc_test( "//xla/hlo/ir:hlo", "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "//xla/runtime:work_dimensions", "//xla/service:buffer_assignment", "//xla/service:buffer_value", "//xla/service:logical_buffer", @@ -134,6 +136,7 @@ xla_cc_test( "//xla/tests:xla_internal_test_main", "//xla/tsl/platform:statusor", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", "@llvm-project//llvm:Support", diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_test.cc b/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_test.cc index 8936c990b56ef7..6ce15a0668363c 100644 --- a/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_test.cc @@ -13,13 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include +#include +#include #include #include "absl/status/statusor.h" #include "absl/strings/string_view.h" +#include "absl/strings/substitute.h" #include "llvm/IR/LLVMContext.h" #include "llvm/Support/raw_ostream.h" #include "mlir/IR/BuiltinOps.h" @@ -29,6 +33,7 @@ limitations under the License. #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" #include "xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h" #include "xla/backends/cpu/codegen/fusion_compiler.h" +#include "xla/backends/cpu/codegen/fusion_emitter.h" #include "xla/codegen/kernel_definition.h" #include "xla/codegen/llvm_kernel_source.h" #include "xla/codegen/mlir_kernel_source.h" @@ -39,6 +44,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/runtime/work_dimensions.h" #include "xla/service/buffer_assignment.h" #include "xla/service/buffer_value.h" #include "xla/service/cpu/cpu_executable.h" @@ -49,7 +55,20 @@ namespace xla { namespace cpu { namespace { -class CpuFusionEmitterTest : public HloHardwareIndependentTestBase { +using ::testing::ElementsAreArray; + +struct GetWorkDimensionsTestCase { + std::string test_name; + absl::string_view shape_and_layout; + int64_t outer_dimension_partitions; + uint64_t expected_num_work_groups; + std::vector expected_work_tile_size; + absl::string_view root_op = "copy(p0)"; +}; + +class CpuFusionEmitterTest + : public HloHardwareIndependentTestBase, + public ::testing::WithParamInterface { protected: absl::StatusOr> RunBufferAssignment( const HloModule& hlo) { @@ -155,6 +174,235 @@ TEST_F(CpuFusionEmitterTest, ScatterLlvm) { EXPECT_TRUE(filecheck_matched); } +TEST_P(CpuFusionEmitterTest, WorkDimensionsTest) { + const GetWorkDimensionsTestCase& tc = GetParam(); + constexpr absl::string_view kHloTemplate = R"( + HloModule test + fused_computation { + p0 = $0 parameter(0) + ROOT res = $0 $2 + } + ENTRY main { + p0 = $0 parameter(0) + ROOT fusion = $0 fusion(p0), kind=kLoop, + calls=fused_computation, + backend_config={"outer_dimension_partitions":["$1"]} + } + )"; + std::string hlo = absl::Substitute(kHloTemplate, tc.shape_and_layout, + tc.outer_dimension_partitions, tc.root_op); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr buffer_assignment, + RunBufferAssignment(*module)); + const auto* fusion = Cast( + module->entry_computation()->root_instruction()); + + auto mlir_context = FusionCompiler::CreateContext(); + ASSERT_OK_AND_ASSIGN( + KernelDefinition kernel_definition, + EmitFusionKernel(*mlir_context, *fusion, buffer_assignment.get(), + /*use_unique_c_name=*/false, + /*enable_tiled_emitter=*/false)); + const WorkDimensions& dims = kernel_definition.spec().work_dimensions(); + EXPECT_EQ(dims.num_work_groups.x, tc.expected_num_work_groups); + EXPECT_THAT(dims.work_tile_size.dimensions, + ElementsAreArray(tc.expected_work_tile_size)); +} + +INSTANTIATE_TEST_SUITE_P( + WorkDimensions, CpuFusionEmitterTest, + ::testing::ValuesIn({ + { + /*test_name=*/"NonSubByteDoesNotRoundOddTile", + /*shape_and_layout=*/"f32[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{3, 3}, + }, + { + // CeilOfRatio(6, 2) = 3 is rounded up to 4 (multiple of 2 + // elements/byte), giving 4 * 3 = 12 elements (6 whole bytes) per + // workgroup. + /*test_name=*/"SubByteS4RoundsOddSplitTileToMultipleOfTwo", + /*shape_and_layout=*/"s4[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{4, 3}, + }, + { + // CeilOfRatio(5, 2) = 3 is also rounded up to 4, so workgroup 0 + // writes 12 elements (6 whole bytes) and workgroup 1 writes the + // remaining 3 elements starting at byte 6. + /*test_name=*/ + "SubByteS4OddRowsStillRoundsSplitTileToMultipleOfTwo", + /*shape_and_layout=*/"s4[5,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{4, 3}, + }, + { + // CeilOfRatio(6, 2) = 3 is rounded up to 4 (multiple of 4 + // elements/byte), giving 4 * 3 = 12 elements (3 whole bytes) per + // workgroup. + /*test_name=*/"SubByteS2RoundsSplitTileToMultipleOfFour", + /*shape_and_layout=*/"s2[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{4, 3}, + }, + { + // CeilOfRatio(12, 2) = 6 is rounded up to 8 (multiple of 8 + // elements/byte), giving 8 * 3 = 24 elements (3 whole bytes) per + // workgroup. + /*test_name=*/"SubByteS1RoundsSplitTileToMultipleOfEight", + /*shape_and_layout=*/"s1[12,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{8, 3}, + }, + { + // PRED is stored 1 element per byte, so odd split tiles are + // not rounded. + /*test_name=*/"PredIsUnpackedAndDoesNotRoundOddTile", + /*shape_and_layout=*/"pred[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{3, 3}, + }, + { + // CeilOfRatio(6, 4) = 2 (already a multiple of 2 elems/byte); + // 3 workgroups of size 2 cover all 6 rows, so num_work_groups is + // shrunk from 4 to 3. + /*test_name=*/"SubByteRecomputesNumWorkGroupsAfterRounding", + /*shape_and_layout=*/"s4[6,3]{1,0}", + /*outer_dimension_partitions=*/4, + /*expected_num_work_groups=*/3, + /*expected_work_tile_size=*/{2, 3}, + }, + { + // Concatenate fusions with packed sub-byte outputs must run in a + // single workgroup covering the entire shape, because concatenated + // operands can start at non-byte-aligned offsets. + /*test_name=*/"ConcatenateSubByteForcesSingleWorkGroup", + /*shape_and_layout=*/"s4[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/1, + /*expected_work_tile_size=*/{6, 3}, + /*root_op=*/"concatenate(p0), dimensions={0}", + }, + { + // DynamicUpdateSlice fusions with packed sub-byte outputs must also + // run in a single workgroup covering the entire shape, because the + // update index may land at a non-byte aligned offset. + /*test_name=*/"DynamicUpdateSliceSubByteForcesSingleWorkGroup", + /*shape_and_layout=*/"s4[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/1, + /*expected_work_tile_size=*/{6, 3}, + /*root_op=*/ + "dynamic-update-slice(p0, p0, s32[] constant(0), s32[] " + "constant(0))", + }, + { + // 3D shape where the leading dim (2) is smaller than 4 workgroups: + // dim 0 (2) is folded with dim 1 (6) into 12 slices of 3 elements; + // CeilOfRatio(12, 4) = 3 rounds up to 4, leaving 3 workgroups of + // tile {4, 3}. + /*test_name=*/"ThreeDimsFoldsLeadingDimAndRoundsSplitTile", + /*shape_and_layout=*/"s4[2,6,3]{2,1,0}", + /*outer_dimension_partitions=*/4, + /*expected_num_work_groups=*/3, + /*expected_work_tile_size=*/{4, 3}, + }, + { + // 3D shape where the leading dim (6) is large enough to split + // directly: CeilOfRatio(6, 2) = 3 rounds up to 4, keeping both + // minor dims (4 and 3) intact. + /*test_name=*/"ThreeDimsSplitsOutermostDimAndKeepsMinorDims", + /*shape_and_layout=*/"s4[6,4,3]{2,1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{4, 4, 3}, + }, + { + // More workgroups (8) than the product of the first two dimensions + // (2 * 2 = 4): all three dimensions fold together (2 * 2 * 3 = 12), + // and CeilOfRatio(12, 8) = 2 yields 6 workgroups of 1D tile {2}. + /*test_name=*/"FoldsAllDimensionsIntoSingleTileDim", + /*shape_and_layout=*/"s4[2,2,3]{2,1,0}", + /*outer_dimension_partitions=*/8, + /*expected_num_work_groups=*/6, + /*expected_work_tile_size=*/{2}, + }, + { + // As many workgroups (12) as total elements (2 * 2 * 3 = 12): all + // three dimensions fold together, and CeilOfRatio(12, 12) = 1 is + // rounded up to 2 elements (1 byte) per workgroup, shrinking + // num_work_groups from 12 to 6. + /*test_name=*/"OneWorkGroupPerElementRoundsToWholeByte", + /*shape_and_layout=*/"s4[2,2,3]{2,1,0}", + /*outer_dimension_partitions=*/12, + /*expected_num_work_groups=*/6, + /*expected_work_tile_size=*/{2}, + }, + { + // Column-major layout {0,1}: dim 1 (size 6) is the most-major + // physical dimension and gets split and rounded from 3 to 4, while + // dim 0 (size 3) is minor. + /*test_name=*/"SubByteColumnMajorLayoutSplitsPhysicalMajorDim", + /*shape_and_layout=*/"s4[3,6]{0,1}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{4, 3}, + }, + { + // Single workgroup does not round odd dimensions. + /*test_name=*/"SingleWorkGroupDoesNotRoundOddSubByteDims", + /*shape_and_layout=*/"s4[5,3]{1,0}", + /*outer_dimension_partitions=*/1, + /*expected_num_work_groups=*/1, + /*expected_work_tile_size=*/{5, 3}, + }, + { + // f6 types are stored one element per byte and are not rounded. + /*test_name=*/"NonPowerOfTwoSubByteF6DoesNotRoundOddTile", + /*shape_and_layout=*/"f6e3m2fn[6,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/2, + /*expected_work_tile_size=*/{3, 3}, + }, + { + // More workgroups (16) than elements (12): clamped to 12, then + // rounded to 2 elements (1 byte) per workgroup. + /*test_name=*/"MoreWorkGroupsThanElementsSubByteRoundsToWholeByte", + /*shape_and_layout=*/"s4[2,2,3]{2,1,0}", + /*outer_dimension_partitions=*/16, + /*expected_num_work_groups=*/6, + /*expected_work_tile_size=*/{2}, + }, + { + // More workgroups (8) than elements (6): clamped to 6. + /*test_name=*/ + "MoreWorkGroupsThanElementsNonSubByteClampsToElements", + /*shape_and_layout=*/"f32[2,3]{1,0}", + /*outer_dimension_partitions=*/8, + /*expected_num_work_groups=*/6, + /*expected_work_tile_size=*/{1}, + }, + { + // Zero-element shapes get a single workgroup and an empty tile. + /*test_name=*/"ZeroElementsUsesSingleWorkGroup", + /*shape_and_layout=*/"f32[0,3]{1,0}", + /*outer_dimension_partitions=*/2, + /*expected_num_work_groups=*/1, + /*expected_work_tile_size=*/{}, + }, + }), + [](const ::testing::TestParamInfo& info) { + return info.param.test_name; + }); + } // namespace } // namespace cpu } // namespace xla diff --git a/third_party/xla/xla/backends/cpu/codegen/execution_engine.cc b/third_party/xla/xla/backends/cpu/codegen/execution_engine.cc index ec5bfbb34536cc..9774f40298ca8a 100644 --- a/third_party/xla/xla/backends/cpu/codegen/execution_engine.cc +++ b/third_party/xla/xla/backends/cpu/codegen/execution_engine.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/log/log.h" #include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" #include "llvm/ExecutionEngine/JITEventListener.h" #include "llvm/ExecutionEngine/Orc/Core.h" #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" @@ -30,6 +31,7 @@ limitations under the License. #include "llvm/Support/MemoryBuffer.h" #include "xla/backends/cpu/codegen/contiguous_section_memory_manager.h" #include "xla/backends/cpu/codegen/jit_memory_mapper.h" +#include "xla/backends/cpu/codegen/object_buffer_identifier.h" namespace xla::cpu { @@ -37,8 +39,10 @@ static std::unique_ptr CreateObjectLinkingLayer(llvm::orc::ExecutionSession& execution_session) { return std::make_unique( execution_session, [](const llvm::MemoryBuffer& obj) { + absl::string_view memory_region_name = + ExtractMemoryRegionName(obj.getBufferIdentifier()); return std::make_unique( - GetJitMemoryMapper(obj.getBufferIdentifier())); + GetJitMemoryMapper(memory_region_name)); }); } diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.cc index f3ed43c8f80d52..2f1946dda170bf 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.cc +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/backends/cpu/codegen/fusion_emitter.h" +#include #include #include #include @@ -26,6 +27,7 @@ limitations under the License. #include "absl/container/inlined_vector.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/numeric/bits.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -55,6 +57,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/utils/hlo_traversal.h" #include "xla/layout_util.h" +#include "xla/primitive_util.h" #include "xla/runtime/work_cluster.h" #include "xla/runtime/work_dimensions.h" #include "xla/runtime/work_group.h" @@ -142,32 +145,80 @@ static int64_t GetWorkGroupCount(const HloFusionInstruction& fusion) { std::multiplies()); } -WorkDimensions GetWorkDimensions(const Shape& shape, - const HloFusionInstruction& fusion) { +// Returns the maximum number of sub-byte elements packed into a single byte +// across all outputs of `shape` (e.g. 2 for 4-bit types, 4 for 2-bit types, or +// 1 if there are no packed sub-byte outputs). +static int64_t GetMaxSubByteElementsPerByte(const Shape& shape) { + int64_t max_elements_per_byte = 1; + ShapeUtil::ForEachSubshape( + shape, [&](const Shape& subshape, const ShapeIndex&) { + if (subshape.IsArray() && + primitive_util::IsSubByteNonPredType(subshape.element_type())) { + int64_t bit_width = primitive_util::BitWidth(subshape.element_type()); + CHECK(absl::has_single_bit(static_cast(bit_width))) + << "Expected power-of-2 sub-byte bit width, got " << bit_width; + max_elements_per_byte = + std::max(max_elements_per_byte, 8 / bit_width); + } + }); + return max_elements_per_byte; +} + +// Partitions the iteration space of `shape` across workgroups along the +// major physical dimensions so each workgroup executes a contiguous slice of +// elements. Leading dimensions smaller than the requested workgroup count are +// folded into the next dimension until there is enough work to split, at which +// point we define the "split tile" size. All remaining minor dimensions are +// kept intact in each tile. +// +// For packed sub-byte outputs, if allow_sub_byte_multi_work_group is: +// - true, the split tile size is rounded up to a multiple of the +// elements-per-byte packing factor, so that distinct workgroups never +// share a byte. +// - false, a single workgroup is used. +static WorkDimensions GetWorkDimensions( + const Shape& shape, const HloFusionInstruction& fusion, + bool allow_sub_byte_multi_work_group = true) { if (!shape.has_layout()) { Shape shape_with_layout = shape; LayoutUtil::SetToDefaultLayout(&shape_with_layout); - return GetWorkDimensions(shape_with_layout, fusion); + return GetWorkDimensions(shape_with_layout, fusion, + allow_sub_byte_multi_work_group); } auto minor_to_major = LayoutUtil::MinorToMajor(shape.layout()); - if (minor_to_major.empty()) { return WorkDimensions{}; } - int64_t work_group_count = GetWorkGroupCount(fusion); + int64_t total_elements = ShapeUtil::ElementsIn(shape); + if (total_elements == 0) { + return WorkDimensions{}; + } + + int64_t sub_byte_alignment = GetMaxSubByteElementsPerByte(fusion.shape()); + bool has_sub_byte_output = sub_byte_alignment > 1; + bool can_parallelize = + allow_sub_byte_multi_work_group || !has_sub_byte_output; + int64_t work_group_count = + can_parallelize ? std::min(GetWorkGroupCount(fusion), total_elements) : 1; NumWorkGroups num_work_groups{static_cast(work_group_count)}; WorkTileSize work_tile_size; int64_t folded_dims = 1; for (int64_t dim : llvm::reverse(minor_to_major)) { int64_t dim_size = ShapeUtil::GetDimension(shape, dim); - int64_t accumilated_dim_size = folded_dims * dim_size; - if (accumilated_dim_size < work_group_count) { + int64_t accumulated_dim_size = folded_dims * dim_size; + if (accumulated_dim_size < work_group_count) { folded_dims *= dim_size; } else if (work_group_count != 1) { - work_tile_size.dimensions.push_back( - CeilOfRatio(accumilated_dim_size, work_group_count)); + // Round `split_tile` up to `sub_byte_alignment` so that + // `split_tile * (product of minor dims)` is a whole number of bytes. + int64_t split_tile = + RoundUpTo(CeilOfRatio(accumulated_dim_size, work_group_count), + sub_byte_alignment); + work_tile_size.dimensions.push_back(split_tile); + num_work_groups = NumWorkGroups{ + static_cast(CeilOfRatio(accumulated_dim_size, split_tile))}; work_group_count = 1; } else { work_tile_size.dimensions.push_back(dim_size); @@ -193,7 +244,8 @@ static WorkDimensions GetConcatenateEmitterWorkDims( Shape indexing_shape = emitters::ConcatenateFusionKernelEmitter::GetIndexingShape(fusion_spec); - return GetWorkDimensions(indexing_shape, fusion); + return GetWorkDimensions(indexing_shape, fusion, + /*allow_sub_byte_multi_work_group=*/false); } static WorkDimensions GetDynamicUpdateSliceEmitterWorkDims( @@ -201,7 +253,8 @@ static WorkDimensions GetDynamicUpdateSliceEmitterWorkDims( Shape indexing_shape = emitters::DynamicUpdateSliceKernelEmitter::GetIndexingShape(fusion_spec); - return GetWorkDimensions(indexing_shape, fusion); + return GetWorkDimensions(indexing_shape, fusion, + /*allow_sub_byte_multi_work_group=*/false); } static HloFusionSpec GetLoopFusionSpec(const HloFusionInstruction& fusion) { diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc index b5ee9661e168ac..acfc960021cb83 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler.cc @@ -68,6 +68,7 @@ limitations under the License. #include "llvm/Transforms/IPO/AlwaysInliner.h" #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h" #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" +#include "xla/backends/cpu/codegen/object_buffer_identifier.h" #include "xla/backends/cpu/codegen/polynomial_approximations.h" #include "xla/backends/cpu/target_machine_options.h" #include "xla/codegen/intrinsic/intrinsic.h" @@ -513,8 +514,17 @@ std::unique_ptr IrCompiler::EmitMachineCode( CHECK(md_str != nullptr); llvm::StringRef mem_region_name_str = md_str->getString(); + // Each module gets assigned two names encoded into the buffer identifier: + // - Memory region name: human-friendly name shared among related kernels, + // so that profilers can aggregate results per kernel. + // - Buffer identifier: to refer to each module uniquely. Necessary for + // sanitizers. + std::string buffer_identifier = EncodeBufferIdentifier( + absl::string_view(mem_region_name_str.data(), mem_region_name_str.size()), + module.getModuleIdentifier()); + return std::make_unique( - std::move(mc_stream_buffer), mem_region_name_str); + std::move(mc_stream_buffer), buffer_identifier); } llvm::CodeGenOptLevel IrCompiler::GetCodeGenOptLevel( diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc index 4f0484caf2dc56..2d48a151a0def5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc +++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc @@ -45,6 +45,7 @@ limitations under the License. #include "llvm/Target/TargetMachine.h" #include "llvm/TargetParser/Triple.h" #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" +#include "xla/backends/cpu/codegen/object_buffer_identifier.h" #include "xla/backends/cpu/target_machine_options.h" #include "xla/debug_options_flags.h" #include "xla/service/cpu/backend_config.pb.h" @@ -93,7 +94,8 @@ constexpr absl::string_view kUnoptimizedIr = R"( // Parses the LLVM IR into a ThreadSafeModule. static absl::StatusOr> ParseModule( - llvm::LLVMContext& context, absl::string_view ir, absl::string_view name) { + llvm::LLVMContext& context, absl::string_view ir, absl::string_view name, + absl::string_view memory_region_name = "ir_compiler_test") { llvm::SMDiagnostic diagnostic; llvm::MemoryBufferRef ir_buffer(ir, name); @@ -103,7 +105,7 @@ static absl::StatusOr> ParseModule( diagnostic.getMessage().str()); } - SetModuleMemoryRegionName(*m, "ir_compiler_test"); + SetModuleMemoryRegionName(*m, memory_region_name); return m; } @@ -434,6 +436,45 @@ INSTANTIATE_TEST_SUITE_P(IrCompilerParameterizedTestInstantiation, ::testing::Values("x86_64-grtev4-linux-gnu", "aarch64-unknown-linux-gnu")); +TEST(IrCompilerTest, EmitsBufferIdentifierWithCoarseAndUniqueNames) { + auto context = std::make_unique(); + std::unique_ptr ir_compiler = + IrCompiler::Create(llvm::TargetOptions(), IrCompiler::Options(), + IrCompiler::CompilationHooks()); + + ASSERT_OK_AND_ASSIGN(auto ir_module, ParseModule(*context, kUnoptimizedIr, + "unique_kernel_module_name", + "coarse_emitter_fusion")); + + ASSERT_OK_AND_ASSIGN(auto target_machine, + ir_compiler->build_target_machine()); + + ir_module->setDataLayout(target_machine->createDataLayout()); + ir_module->setTargetTriple(target_machine->getTargetTriple()); + + auto memory_buffer = cantFail((*ir_compiler)(*ir_module)); + ASSERT_NE(memory_buffer, nullptr); + + absl::string_view buffer_identifier = memory_buffer->getBufferIdentifier(); + EXPECT_EQ(ExtractMemoryRegionName(buffer_identifier), + "coarse_emitter_fusion"); + EXPECT_EQ(ExtractModuleIdentifier(buffer_identifier), + "unique_kernel_module_name"); +} + +TEST(ObjectBufferIdentifierTest, EncodeAndExtract) { + std::string encoded = + EncodeBufferIdentifier("my_region", "unique_module_123"); + EXPECT_EQ(ExtractMemoryRegionName(encoded), "my_region"); + EXPECT_EQ(ExtractModuleIdentifier(encoded), "unique_module_123"); + + // Backward compatibility with raw / un-delimited identifiers. + EXPECT_EQ(ExtractMemoryRegionName("legacy_region"), "legacy_region"); + EXPECT_EQ(ExtractModuleIdentifier("legacy_region"), "legacy_region"); + EXPECT_EQ(ExtractMemoryRegionName(""), ""); + EXPECT_EQ(ExtractModuleIdentifier(""), ""); +} + } // namespace } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/object_buffer_identifier.h b/third_party/xla/xla/backends/cpu/codegen/object_buffer_identifier.h new file mode 100644 index 00000000000000..7ae77f0cb09037 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/object_buffer_identifier.h @@ -0,0 +1,65 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_CPU_CODEGEN_OBJECT_BUFFER_IDENTIFIER_H_ +#define XLA_BACKENDS_CPU_CODEGEN_OBJECT_BUFFER_IDENTIFIER_H_ + +#include +#include + +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" + +namespace xla::cpu { + +// Delimiter that separates the memory region name (used for profiling) from the +// unique module identifier (used for sanitizers and symbol identification). +inline constexpr absl::string_view kMemoryRegionSeparator = + "::__xla_mem_region_separator__::"; + +// Encodes the memory region name and module identifier into a single object +// buffer identifier. +inline std::string EncodeBufferIdentifier(absl::string_view memory_region_name, + absl::string_view module_identifier) { + return absl::StrCat(memory_region_name, kMemoryRegionSeparator, + module_identifier); +} + +// Extracts the memory region name from an object buffer identifier. If the +// separator is found, returns the region name prefix. Otherwise, returns the +// identifier itself for backward compatibility with legacy or raw buffer +// identifiers. +inline absl::string_view ExtractMemoryRegionName( + absl::string_view buffer_identifier) { + size_t pos = buffer_identifier.find(kMemoryRegionSeparator); + if (pos != absl::string_view::npos) { + return buffer_identifier.substr(0, pos); + } + return buffer_identifier; +} + +// Extracts the unique module identifier from an object buffer identifier. +inline absl::string_view ExtractModuleIdentifier( + absl::string_view buffer_identifier) { + size_t pos = buffer_identifier.find(kMemoryRegionSeparator); + if (pos != absl::string_view::npos) { + return buffer_identifier.substr(pos + kMemoryRegionSeparator.size()); + } + return buffer_identifier; +} + +} // namespace xla::cpu + +#endif // XLA_BACKENDS_CPU_CODEGEN_OBJECT_BUFFER_IDENTIFIER_H_ diff --git a/third_party/xla/xla/backends/cpu/testlib/BUILD b/third_party/xla/xla/backends/cpu/testlib/BUILD index 64deb072a7a7c1..92ed72a1470048 100644 --- a/third_party/xla/xla/backends/cpu/testlib/BUILD +++ b/third_party/xla/xla/backends/cpu/testlib/BUILD @@ -175,7 +175,6 @@ tsl_pybind_extension( "//xla/backends/cpu/codegen:target_machine_features", "//xla/backends/cpu/codegen/dot:dot_kernel_emitter", "//xla/backends/cpu/codegen/elemental:concatenate_kernel_emitter", - "//xla/backends/cpu/codegen/elemental:elemental_kernel_emitter", "//xla/backends/cpu/codegen/emitters:cpu_fusion_emitters", "//xla/codegen:kernel_definition", "//xla/codegen:kernel_emitter", @@ -239,21 +238,3 @@ py_test( "@absl_py//absl/testing:absltest", ], ) - -py_test( - name = "elemental_kernel_emitter_test", - srcs = ["elemental_kernel_emitter_test.py"], - main = "elemental_kernel_emitter_test.py", - strict_deps = True, - tags = [ - "no_oss", - ], - deps = [ - ":testlib", - "//third_party/py/numpy", - "//xla/codegen/testlib", - "//xla/python:xla_extension", - "@absl_py//absl/testing:absltest", - "@absl_py//absl/testing:parameterized", - ], -) diff --git a/third_party/xla/xla/backends/cpu/testlib/__init__.py b/third_party/xla/xla/backends/cpu/testlib/__init__.py index 3261fc946d3f0a..4cb6cf5999a357 100644 --- a/third_party/xla/xla/backends/cpu/testlib/__init__.py +++ b/third_party/xla/xla/backends/cpu/testlib/__init__.py @@ -21,7 +21,6 @@ ComputationKernelEmitter = _extension.ComputationKernelEmitter ConcatenateKernelEmitter = _extension.ConcatenateKernelEmitter DotKernelEmitter = _extension.DotKernelEmitter -ElementalKernelEmitter = _extension.ElementalKernelEmitter HloCompiler = _extension.HloCompiler JitCompiler = _extension.JitCompiler KernelRunner = _extension.KernelRunner diff --git a/third_party/xla/xla/backends/cpu/testlib/elemental_kernel_emitter_test.py b/third_party/xla/xla/backends/cpu/testlib/elemental_kernel_emitter_test.py deleted file mode 100644 index 298dcbe98810a6..00000000000000 --- a/third_party/xla/xla/backends/cpu/testlib/elemental_kernel_emitter_test.py +++ /dev/null @@ -1,391 +0,0 @@ -# Copyright 2024 The OpenXLA Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -from collections.abc import Callable, Sequence -import dataclasses -import itertools -import math - -from absl.testing import absltest -from absl.testing import parameterized -import numpy as np - -from xla.backends.cpu import testlib as testlib_cpu -from xla.backends.cpu.testlib import utilities -from xla.codegen import testlib as testlib_base -from xla.codegen.testlib.utilities import opcode_arity -from xla.python import xla_extension - -HloOpcode = testlib_base.HloOpcode -create_literal = testlib_base.utilities.create_literal_from_np -HloInstruction = testlib_base.HloInstruction -ComparisonDirection = testlib_base.ComparisonDirection -_inf = float("inf") - - -def create_input( - value_range: tuple[float, float], - shape: Sequence[int], - dtype: np.dtype, - shuffle: bool = False, -) -> np.ndarray: - size = np.prod(shape) if shape else 1 - result = np.linspace( - value_range[0], value_range[1], size, dtype=dtype - ).reshape(shape) - - if shuffle and (np.ndim(result) != 0): - np.random.shuffle(result) - - return result - - -def np_erf(x): - return np.vectorize(math.erf, otypes=[x.dtype])(x) - - -@dataclasses.dataclass(frozen=True) -class ElementalHloOpcodeDef: - op: HloOpcode - np_op: Callable[[np.ndarray, ...], np.ndarray] - input_ranges: tuple[float, float] = (-1.0, 1.0) - decimal_precision: int = 6 - - # For simple unpacking - def __iter__(self): - return iter( - (self.op, self.np_op, self.input_ranges, self.decimal_precision) - ) - - def __repr__(self): - return f"{self.op.name}({self.input_ranges})" - - -@parameterized.product( - op_def=[ - ElementalHloOpcodeDef(HloOpcode.sine, np.sin), - ElementalHloOpcodeDef(HloOpcode.cosine, np.cos), - ElementalHloOpcodeDef(HloOpcode.tan, np.tan), - ElementalHloOpcodeDef(HloOpcode.exponential, np.exp), - ElementalHloOpcodeDef(HloOpcode.log, np.log, (0.01, 10.0)), - ElementalHloOpcodeDef(HloOpcode.log_plus_one, np.log1p), - ElementalHloOpcodeDef(HloOpcode.sqrt, np.sqrt), - ElementalHloOpcodeDef( - HloOpcode.rsqrt, lambda x: np.reciprocal(np.sqrt(x)) - ), - ElementalHloOpcodeDef(HloOpcode.cbrt, np.cbrt), - ElementalHloOpcodeDef(HloOpcode.power, np.pow), - ElementalHloOpcodeDef(HloOpcode.add, np.add), - ElementalHloOpcodeDef(HloOpcode.subtract, np.subtract), - ElementalHloOpcodeDef(HloOpcode.multiply, np.multiply), - ElementalHloOpcodeDef(HloOpcode.divide, np.divide), - ElementalHloOpcodeDef(HloOpcode.maximum, np.maximum), - ElementalHloOpcodeDef(HloOpcode.minimum, np.minimum), - ElementalHloOpcodeDef(HloOpcode.sign, np.sign), - ElementalHloOpcodeDef(HloOpcode.negate, np.negative), - ElementalHloOpcodeDef(HloOpcode.is_finite, np.isfinite, (-_inf, _inf)), - ElementalHloOpcodeDef(HloOpcode.ceil, np.ceil, (-10.0, 10.0)), - ElementalHloOpcodeDef(HloOpcode.floor, np.floor, (-5.0, 5.0)), - ElementalHloOpcodeDef(HloOpcode.tanh, np.tanh), - ElementalHloOpcodeDef(HloOpcode.atan2, np.arctan2), - ElementalHloOpcodeDef(HloOpcode.erf, np_erf), - ElementalHloOpcodeDef(HloOpcode.exponential_minus_one, np.expm1), - # TODO(willfroom): Update to use better inputs for the following. - ElementalHloOpcodeDef(HloOpcode.clamp, np.clip), - # TODO(willfroom): Add complex ops. - # ElementalHloOpcodeDef(HloOpcode.complex, np.complex), - # ElementalHloOpcodeDef(HloOpcode.real, np.real), - # ElementalHloOpcodeDef(HloOpcode.imag, np.imag), - # TODO(willfroom): go through ElementalIrEmitter interface and ensure - # that all ops are implemented. - # ... - ], - shape=[(4,), (4, 3), (4, 3, 10)], - dtype=[np.dtype(np.float32), np.dtype(np.float64)], -) -class ElementalKernelRunnerTest(parameterized.TestCase): - - def test_elemental_kernel_emitter( - self, - op_def: ElementalHloOpcodeDef, - shape: tuple[int, ...], - dtype: np.dtype, - ): - - [op, np_op, input_ranges, decimal_precision] = op_def - - num_inputs = opcode_arity(op) - self.assertIsNotNone(num_inputs) - - np_inputs = [ - create_input(input_ranges, shape, dtype) for _ in range(num_inputs) - ] - input_literals = [create_literal(input_array) for input_array in np_inputs] - - expected_output = np_op(*np_inputs) - output_literal = create_literal( - np.ndarray(shape, dtype=expected_output.dtype) - ) - - hlo_parameters = [ - HloInstruction.create_parameter(idx, literal.shape(), f"input_{idx}") - for [idx, literal] in enumerate(input_literals) - ] - - hlo_op = HloInstruction.create_variadic( - output_literal.shape(), op, hlo_parameters - ) - - hlo_module, buffer_assignment = utilities.build_hlo_module( - hlo_op, *hlo_parameters - ) - jit_compiler = testlib_cpu.JitCompiler(hlo_module.get_config()) - - emitter = testlib_cpu.ElementalKernelEmitter( - hlo_module.get_root_instruction(), - buffer_assignment, - jit_compiler.get_target_machine(), - ) - - kernel_definition = emitter.emit_kernel_definition() - self.assertIsNotNone(kernel_definition) - - # kernel_definition is consumed by the runner, so we need to save the IR - # string before passing it to the runner. - ir_string = str(kernel_definition.source()) - - runner = testlib_cpu.KernelRunner.create(kernel_definition, jit_compiler) - - runner.call(list(itertools.chain(input_literals, [output_literal]))) - np.testing.assert_array_almost_equal( - np.asarray(output_literal), - expected_output, - decimal=decimal_precision, - err_msg=ir_string, - ) - - -@parameterized.product( - op_def=[ - (ComparisonDirection.kEq, np.equal), - (ComparisonDirection.kNe, np.not_equal), - (ComparisonDirection.kGe, np.greater_equal), - (ComparisonDirection.kGt, np.greater), - (ComparisonDirection.kLe, np.less_equal), - (ComparisonDirection.kLt, np.less), - ], - shape=[(4,), (4, 3), (4, 3, 10)], - dtype=[ - np.dtype(np.uint8), - np.dtype(np.uint16), - np.dtype(np.uint32), - np.dtype(np.uint64), - np.dtype(np.int8), - np.dtype(np.int16), - np.dtype(np.int32), - np.dtype(np.int64), - np.dtype(np.float16), - np.dtype(np.float32), - np.dtype(np.float64), - ], -) -class ElementalComparisonKernelRunnerTest(parameterized.TestCase): - - def test_elemental_comparision_kernel_emitter(self, op_def, shape, dtype): - [direction, np_op] = op_def - - is_unsigned = np.issubdtype(dtype, np.unsignedinteger) - value_range = (0.0, 20.0) if is_unsigned else (-10.0, 10.0) - lhs_np = create_input(value_range, shape, dtype, shuffle=True) - rhs_np = create_input(value_range, shape, dtype, shuffle=True) - - lhs_literal = create_literal(lhs_np) - rhs_literal = create_literal(rhs_np) - - output_literal = create_literal(np.ndarray(shape, dtype=np.bool)) - - lhs_param = HloInstruction.create_parameter(0, lhs_literal.shape(), "lhs") - rhs_param = HloInstruction.create_parameter(1, rhs_literal.shape(), "rhs") - - hlo_op = HloInstruction.create_compare( - output_literal.shape(), lhs_param, rhs_param, direction - ) - - hlo_module, buffer_assignment = utilities.build_hlo_module( - hlo_op, lhs_param, rhs_param - ) - jit_compiler = testlib_cpu.JitCompiler(hlo_module.get_config()) - - emitter = testlib_cpu.ElementalKernelEmitter( - hlo_module.get_root_instruction(), - buffer_assignment, - jit_compiler.get_target_machine(), - ) - - runner = testlib_cpu.KernelRunner.create( - emitter.emit_kernel_definition(), jit_compiler - ) - - runner.call([lhs_literal, rhs_literal, output_literal]) - np.testing.assert_equal( - np.asarray(output_literal), - np_op(lhs_np, rhs_np), - ) - - -@parameterized.product( - input_dimensions=[(4,), (4, 3), (4, 3, 10)], - dtype=[ - np.dtype(np.uint8), - np.dtype(np.uint16), - np.dtype(np.uint32), - np.dtype(np.uint64), - np.dtype(np.int8), - np.dtype(np.int16), - np.dtype(np.int32), - np.dtype(np.int64), - np.dtype(np.float16), - np.dtype(np.float32), - np.dtype(np.float64), - ], -) -class HloModuleKernelRunnerTest(parameterized.TestCase): - - def test_map(self, input_dimensions, dtype): - scalar_shape = xla_extension.Shape.scalar_shape(dtype) - shape = xla_extension.Shape.array_shape(dtype, input_dimensions) - - # Please note the double curly braces is to escape the python string - # formatting. - hlo = """ - HloModule test_map - - double {{ - a = {scalar_shape} parameter(0) - b = {scalar_shape} constant(2) - ROOT doubled = {scalar_shape} multiply(a, b) - }} - - ENTRY main {{ - a = {shape} parameter(0) - ROOT mapped = {shape} map(a), to_apply=double - }} - """.format(scalar_shape=scalar_shape, shape=shape) - - hlo_module, buffer_assignment = utilities.parse_hlo_module(hlo) - jit_compiler = testlib_cpu.JitCompiler(hlo_module.get_config()) - - emitter = testlib_cpu.ElementalKernelEmitter( - hlo_module.get_root_instruction(), - buffer_assignment, - jit_compiler.get_target_machine(), - ) - - input_np = create_input([0, 10], input_dimensions, dtype, shuffle=True) - - input_literal = create_literal(input_np) - - output_literal = xla_extension.Literal(shape) - - runner = testlib_cpu.KernelRunner.create( - emitter.emit_kernel_definition(), jit_compiler - ) - - runner.call([input_literal, output_literal]) - - np.testing.assert_equal( - np.asarray(output_literal), - input_np * 2, - ) - - def test_reduce(self, input_dimensions, dtype): - # Iterate over all combinations of reduce dimensions. - for reduce_dimensions in itertools.chain.from_iterable( - itertools.combinations(range(len(input_dimensions)), r) - for r in range(1, len(input_dimensions)) - ): - scalar_shape = xla_extension.Shape.scalar_shape(dtype) - input_shape = xla_extension.Shape.array_shape(dtype, input_dimensions) - - output_dimensions = [ - dim - for idx, dim in enumerate(input_dimensions) - if idx not in reduce_dimensions - ] - # Result can overflow in int8 (which results in undefined behavior), - # so we use int16 instead. - output_dtype = np.dtype(np.int16) if (dtype == np.int8) else dtype - output_shape = xla_extension.Shape.array_shape( - output_dtype, output_dimensions - ) - - # Please note the double curly braces is to escape the python string - # formatting. - hlo = """ - HloModule test_reduce - - add_method {{ - a = {scalar_shape} parameter(0) - b = {scalar_shape} parameter(1) - ROOT add = {scalar_shape} add(a, b) - }} - - ENTRY main {{ - array = {input_shape} parameter(0) - initial_value = {scalar_shape} parameter(1) - ROOT reduced = {output_shape} reduce(array, initial_value), - dimensions={{{reduce_dimensions}}}, to_apply=add_method - }} - """.format( - scalar_shape=scalar_shape, - input_shape=input_shape, - reduce_dimensions=",".join(map(str, reduce_dimensions)), - output_shape=output_shape, - ) - - hlo_module, buffer_assignment = utilities.parse_hlo_module(hlo) - jit_compiler = testlib_cpu.JitCompiler(hlo_module.get_config()) - - emitter = testlib_cpu.ElementalKernelEmitter( - hlo_module.get_root_instruction(), - buffer_assignment, - jit_compiler.get_target_machine(), - ) - - input_np = create_input([0, 10], input_dimensions, dtype) - input_literal = create_literal(input_np) - - initial_value_np = create_input([0, 10], (), dtype) - initial_value_literal = create_literal(initial_value_np) - - output_literal = xla_extension.Literal(output_shape) - - runner = testlib_cpu.KernelRunner.create( - emitter.emit_kernel_definition(), jit_compiler - ) - - runner.call([input_literal, initial_value_literal, output_literal]) - - np.testing.assert_array_almost_equal_nulp( - np.asarray(output_literal), - np.add.reduce( - input_np, axis=reduce_dimensions, initial=initial_value_np - ), - nulp=3, - ) - - -if __name__ == "__main__": - absltest.main() diff --git a/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc b/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc index 2babfc1dec6f42..5d433df4bad9ca 100644 --- a/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc +++ b/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc @@ -32,7 +32,6 @@ limitations under the License. #include "xla/backends/cpu/codegen/computation_kernel_emitter.h" #include "xla/backends/cpu/codegen/dot/dot_kernel_emitter.h" #include "xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.h" -#include "xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h" #include "xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h" #include "xla/backends/cpu/codegen/fusion_compiler.h" #include "xla/backends/cpu/codegen/fusion_emitter.h" @@ -156,13 +155,6 @@ NB_MODULE(_extension, kernel_runner_module) { "TargetMachineFeatures") .def("__str__", &TargetMachineFeatures::get_target_feature_string); - nb::class_>( - kernel_runner_module, "ElementalKernelEmitter") - .def(nb::init(), - nb::keep_alive<1, 2>(), nb::keep_alive<1, 3>(), - nb::keep_alive<1, 4>()); - nb::class_>( kernel_runner_module, "DotKernelEmitter") .def(nb::init> hlo_module, const TargetMachineFeatures* target_machine_features) { - FusionWrapper fusion_wrapper( - /*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/true, target_machine_features); + FusionWrapper fusion_wrapper(target_machine_features); absl::StatusOr result = fusion_wrapper.Run(hlo_module.get()); if (!result.ok()) { throw std::runtime_error(std::string(result.status().message())); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc index d1ce98bb6c65f2..f44abb7fc7ccf7 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc @@ -372,6 +372,56 @@ mlir::LogicalResult PopulateReductionComputation( return mlir::success(); } +// Emits code that reads the last barrier signal value posted by this block from +// its own slot in the local rank's signal buffer (`SignalBuffers[rank] +// [block_id * world_size + rank]`) and returns `last + 1` as the signal value +// for this launch's first barrier. It replaces the host-provided invocation +// count argument (see `CollectiveCodegenConfig::device_sync_count`). +// +// `BlockBarrierOp` writes the barrier signal value into +// `SignalBuffers[0..world_size][block_id * world_size + rank]` (after a local +// CTA barrier in `AllReduce`, or from warp 0 in `AllGather`), and no other rank +// or block ever writes to `SignalBuffers[rank][block_id * world_size + rank]`. +// For two-shot `AllReduce`, the second barrier writes `signal_value + 1`, so +// the slot advances by 2 across launches and adding 1 at entry yields the next +// launch's first signal value. Signal buffers are allocated and zeroed once per +// executor, so the counter persists across executions (including CUDA graph +// replays and device-side loops) without any extra store or thread barrier. +mlir::Value EmitDeviceInvocationCount(mlir::ImplicitLocOpBuilder& b, + mlir::Value signal_buffers, + mlir::Value rank, mlir::Value block_id, + int64_t world_size) { + const auto ptr_to_i32_type = + ttir::PointerType::get(b.getI32Type(), kGlobalAddressSpace); + const auto ptr_to_i64_type = + ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace); + mlir::Value signal_buffers_i64 = + ttir::BitcastOp::create(b, ptr_to_i64_type, signal_buffers); + // SignalBuffers[rank] + mlir::Value local_signal_buffer_i64 = ttir::LoadOp::create( + b, ttir::AddPtrOp::create(b, ptr_to_i64_type, signal_buffers_i64, rank), + ttir::CacheModifier::NONE, ttir::EvictionPolicy::NORMAL, + /*isVolatile=*/false); + mlir::Value local_signal_buffer = + ttir::IntToPtrOp::create(b, ptr_to_i32_type, local_signal_buffer_i64); + // SignalBuffers[rank][block_id * world_size + rank] + mlir::Value counter_index = arith::AddIOp::create( + b, + arith::MulIOp::create( + b, block_id, + arith::ConstantOp::create( + b, b.getI32IntegerAttr(static_cast(world_size)))), + rank); + mlir::Value counter_ptr = ttir::AddPtrOp::create( + b, ptr_to_i32_type, local_signal_buffer, counter_index); + // Volatile to make sure that every launch reads the counter from memory. + mlir::Value counter = ttir::LoadOp::create( + b, counter_ptr, ttir::CacheModifier::NONE, ttir::EvictionPolicy::NORMAL, + /*isVolatile=*/true); + return arith::AddIOp::create( + b, counter, arith::ConstantOp::create(b, b.getI32IntegerAttr(1))); +} + class AllReduceEmitter { public: static mlir::LogicalResult Emit(AllReduceEmitterContext ctx, @@ -414,8 +464,8 @@ class AllReduceEmitter { const int32_t start_idx = ctx_.num_input_output_args; device_rank_ = ctx_.xtile_entry_fn.getArgument(start_idx); CHECK(device_rank_.getType().isInteger(32)); - signal_value_ = ctx_.xtile_entry_fn.getArgument(start_idx + 1); - CHECK(signal_value_.getType().isInteger(32)); + // The invocation count argument (`start_idx + 1`) is unused: the signal + // value comes from a counter in device memory (see Emit setup IR below). // !tt.ptr signal_buffers_ = ctx_.xtile_entry_fn.getArgument(start_idx + 2); // !tt.ptr @@ -460,6 +510,9 @@ class AllReduceEmitter { << "Subtile shape: " << absl::StrJoin(subtile_shape_, ","); } // 3. Emit setup IR. + mlir::Value block_id = ttir::GetProgramIdOp::create(builder_, 0); + signal_value_ = EmitDeviceInvocationCount( + builder_, signal_buffers_, device_rank_, block_id, ctx_.world_size); if (remote_input_buffers_.getType() == ptr_to_i64_type_) { remote_input_buffers_i64_ = remote_input_buffers_; } else { @@ -1246,9 +1299,9 @@ absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, int32_t total_args = entry_func.getNumArguments(); int32_t opaque_start = total_args - num_opaque_attr.getInt() - kNumTileIndexArgs; - // Layout: opaque[0]=rank, opaque[1]=signal_value, opaque[2]=signal_buffers + // Layout: opaque[0]=rank, opaque[1]=invocation count (unused, the signal + // value comes from a counter in device memory), opaque[2]=signal_buffers. mlir::Value rank_arg = entry_func.getArgument(opaque_start); - mlir::Value signal_value_arg = entry_func.getArgument(opaque_start + 1); mlir::Value signal_buffers_arg = entry_func.getArgument(opaque_start + 2); // Insert at the beginning of the entry block, right after program_id @@ -1258,10 +1311,13 @@ absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, auto loc = entry_func.getLoc(); mlir::ImplicitLocOpBuilder builder(loc, &entry_block, entry_block.begin()); + mlir::Value block_id = ttir::GetProgramIdOp::create(builder, 0); + mlir::Value signal_value = EmitDeviceInvocationCount( + builder, signal_buffers_arg, rank_arg, block_id, world_size); // Inter-block barrier via signal flags. This blocks until all // remote ranks have also signaled. mtx::BlockBarrierOp::create(builder, signal_buffers_arg, rank_arg, - signal_value_arg, + signal_value, builder.getI32IntegerAttr(world_size)); return absl::OkStatus(); diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc b/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc index 6fae64d06bfc52..0e9a6751a3de85 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/lowering_util.cc @@ -356,6 +356,28 @@ std::pair CreateTensorOfPointersAndMask( range_tile = add_if(arith::AddIOp(), range, range_tile); } + mlir::Value reduced_mask; + for (unsigned dim : reduced_dims) { + if (!IsGuaranteedInBounds(offsets[dim], sizes[dim], original_shape[dim])) { + mlir::Value upper_bound = + arith::ConstantIntOp::create(builder, i64_type, original_shape[dim]); + mlir::Value mask_right = arith::CmpIOp::create( + builder, arith::CmpIPredicate::slt, cast_offsets[dim], upper_bound); + + mlir::Value lower_bound = + arith::ConstantIntOp::create(builder, i64_type, 0); + mlir::Value mask_left = arith::CmpIOp::create( + builder, arith::CmpIPredicate::sge, cast_offsets[dim], lower_bound); + mlir::Value mask = arith::AndIOp::create(builder, mask_left, mask_right); + reduced_mask = add_if(arith::AndIOp(), mask, reduced_mask); + } + } + if (reduced_mask) { + reduced_mask = ttir::SplatOp::create( + builder, i64_tile_type.clone(builder.getI1Type()), reduced_mask); + mask_tile = add_if(arith::AndIOp(), reduced_mask, mask_tile); + } + // Sum up block-uniform offsets multiplied by strides. mlir::Value block_offset; for (auto [cast_offset, shape_stride] : diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD index 1532db909a751e..22cfd08b2cd3b6 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD @@ -68,6 +68,7 @@ xla_cc_binary( deviceless_tests = [ "collectives/all_reduce.hlo", "collectives/all_gather.hlo", + "collectives/all_gather_one_shot.hlo", "concat/concat_active_scalar_constant.hlo", "concat/concat_offset_not_a_multiple_of_tile_size.hlo", "constraints/max_tensor_elements_violation.hlo", diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo new file mode 100644 index 00000000000000..f04d9c1b766fe6 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather_one_shot.hlo @@ -0,0 +1,58 @@ +// RUN: fusion_to_triton %s | FileCheck %s +// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s +f { + %param0 = f32[128,128]{1,0} parameter(0) + ROOT result = f32[256,128]{1,0} all-gather(%param0), + replica_groups={{0,1}}, dimensions={0}, + backend_config={"collective_backend_config":{ + "kernel_strategy":"KERNEL_STRATEGY_TRITON_ONE_SHOT"}} +} + +ENTRY entry { + p0 = f32[128,128]{1,0} parameter(0) + ROOT %fusion = f32[256,128]{1,0} fusion(p0), + kind=kCustom, calls=f, + backend_config={ + "fusion_backend_config": { + "kind": "__triton_collective", + "block_level_fusion_config": { + "num_warps": "4", + "output_tiles": [{sizes: [16,16]}], + "num_ctas": 1, + "num_stages": 1, + "is_tma_allowed": false, + "is_warp_specialization_allowed": false + } + } + } +} + +// One-shot all-gather waits on an entry barrier before reading the peers' +// scratch buffers. The signal value of the barrier is derived from this block's +// own slot in the local signal buffer (`SignalBuffers[rank] +// [block_id * world_size + rank]`). + +// CHECK: xtile.entry_func @triton_fn( +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<2xi64>, +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<256x128xf32>, +// CHECK-SAME: %[[RANK:[a-zA-Z0-9_]+]]: i32, +// CHECK-SAME: %[[INVOCATION_COUNT:[a-zA-Z0-9_]+]]: i32, +// CHECK-SAME: %[[SIGNAL_BUFFERS:[a-zA-Z0-9_]+]]: !tt.ptr, +// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index{{( \{xla.range = \[0 : index, 127 : index\]\})?}}) +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 +// CHECK-DAG: %[[WORLD_SIZE:.*]] = arith.constant 2 : i32 +// CHECK: %[[BLOCK_ID:.*]] = tt.get_program_id x +// CHECK: %[[SIGNAL_BUFFER_ADDR:.*]] = tt.addptr %[[SIGNAL_BUFFERS]], %[[RANK]] +// CHECK: %[[SIGNAL_BUFFER_INT:.*]] = tt.load %[[SIGNAL_BUFFER_ADDR]] +// CHECK: %[[SIGNAL_BUFFER:.*]] = tt.int_to_ptr %[[SIGNAL_BUFFER_INT]] +// CHECK: %[[BLOCK_OFFSET:.*]] = arith.muli %[[BLOCK_ID]], %[[WORLD_SIZE]] +// CHECK: %[[COUNTER_IDX:.*]] = arith.addi %[[BLOCK_OFFSET]], %[[RANK]] +// CHECK: %[[COUNTER_PTR:.*]] = tt.addptr %[[SIGNAL_BUFFER]], %[[COUNTER_IDX]] +// CHECK: %[[COUNTER:.*]] = tt.load %[[COUNTER_PTR]] {isVolatile = true} +// CHECK: %[[SIGNAL_VALUE:.*]] = arith.addi %[[COUNTER]], %[[C1]] +// CHECK: triton_xla.block_barrier %[[SIGNAL_BUFFERS]], %[[RANK]], %[[SIGNAL_VALUE]], +// CHECK-SAME: {world_size = 2 : i32} +// CHECK: xtile.select_buffer %[[ARG0]] +// CHECK: xtile.extract +// CHECK: xtile.insert +// CHECK: xtile.return diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo index 66986fdad786db..8ae71f7f0d3ba6 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo @@ -21,15 +21,29 @@ fusion{ // CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: !tt.ptr, // CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index{{( \{xla.range = \[0 : index, 15 : index\]\})?}}) // CHECK-SAME: attributes {num_opaque_args = 4 : i32} { +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32 +// CHECK-DAG: %[[WORLD_SIZE:.*]] = arith.constant 2 : i32 // CHECK: %[[EXTRACT_0:.*]] = xtile.extract // CHECK-SAME: : memref<65536xf32> -> tensor<4096xf32> +// The signal value is derived from this block's own slot in the local signal +// buffer (`SignalBuffers[rank][block_id * world_size + rank]`). +// CHECK: %[[BLOCK_ID:.*]] = tt.get_program_id x +// CHECK: %[[SIGNAL_BUFFER_ADDR:.*]] = tt.addptr %[[ARG4]], %[[ARG2]] +// CHECK: %[[SIGNAL_BUFFER_INT:.*]] = tt.load %[[SIGNAL_BUFFER_ADDR]] +// CHECK: %[[SIGNAL_BUFFER:.*]] = tt.int_to_ptr %[[SIGNAL_BUFFER_INT]] +// CHECK: %[[BLOCK_OFFSET:.*]] = arith.muli %[[BLOCK_ID]], %[[WORLD_SIZE]] +// CHECK: %[[COUNTER_IDX:.*]] = arith.addi %[[BLOCK_OFFSET]], %[[ARG2]] +// CHECK: %[[COUNTER_PTR:.*]] = tt.addptr %[[SIGNAL_BUFFER]], %[[COUNTER_IDX]] +// CHECK: %[[COUNTER:.*]] = tt.load %[[COUNTER_PTR]] {isVolatile = true} +// CHECK: %[[SIGNAL_VALUE:.*]] = arith.addi %[[COUNTER]], %[[C1]] + // CHECK: tt.store %{{[a-zA-Z0-9_]+}}, %{{[a-zA-Z0-9_]+}} // CHECK-SAME: : tensor<4096x!tt.ptr> // CHECK: ttg.barrier local -// CHECK: triton_xla.block_barrier +// CHECK: triton_xla.block_barrier %[[ARG4]], %[[ARG2]], %[[SIGNAL_VALUE]], // CHECK-SAME: {world_size = 2 : i32} : (!tt.ptr, i32, i32) // CHECK: xtile.insert diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo index 425848af2d620c..8d4a75e4a2b9ab 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo @@ -1,7 +1,5 @@ -// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s -// RUN: not triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s - -// CHECK: MapSymbolIdToSequentialDimValue failed for symbolic dimension +// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s +// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s max_region { p0 = f32[] parameter(0) @@ -40,3 +38,22 @@ ENTRY main { "num_ctas":"1", "num_stages":"1"}}} } + +// `max` feeds both the root (via `exp`) and the tiled `sum`, so tiling +// duplicates it into the region of `sum`. +// TODO(b/559130830): Hoist the loop-invariant `max` out of the `sum` region. + +// Loop 1: top-level `max`, used by the numerator exp(p - max). +// CHECK-LABEL: xtile.entry_func @triton_fn +// CHECK: scf.for +// CHECK: arith.maximumf +// CHECK: %[[EXP:.*]] = tt.extern_elementwise {{.*}} symbol = "__nv_expf" + +// Loop 2: `sum`, with loop 3 (the duplicated `max`) nested inside. +// CHECK: %[[SUM:.*]] = scf.for {{.*}} iter_args(%[[ACC:[^ ]+]] = +// CHECK: scf.for +// CHECK: arith.maximumf +// CHECK: %[[INNER_EXP:.*]] = tt.extern_elementwise {{.*}} symbol = "__nv_expf" +// CHECK: arith.addf %[[ACC]], %[[INNER_EXP]] +// CHECK: "tt.reduce"(%[[SUM]]) +// CHECK: arith.divf %[[EXP]], diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tools/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tools/BUILD index a028e37eab25ad..0feec115b66e4e 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tools/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tools/BUILD @@ -57,3 +57,38 @@ xla_cc_binary( "@tsl//tsl/platform:platform_port", ], ) + +xla_cc_binary( + name = "fusion_to_tiled_hlo", + srcs = ["fusion_to_tiled_hlo.cc"], + linkopts = ["-Wl,-rpath,$$ORIGIN/../lit_lib"], + deps = [ + "//xla:debug_options_flags", + "//xla:xla_proto_cc", + "//xla/codegen/tiling/experimental:tile", + "//xla/codegen/tiling/experimental:tiled_hlo", + "//xla/codegen/xtile:block_level_parameters", + "//xla/codegen/xtile:tiling_from_block_parameters", + "//xla/hlo/analysis:symbolic_map", + "//xla/hlo/ir:hlo", + "//xla/hlo/utils:hlo_traversal", + "//xla/service/gpu:backend_configs_cc", + "//xla/tools:hlo_module_loader", + "//xla/tsl/util:command_line_flags", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@tsl//tsl/platform:platform_port", + ], +) + +alias( + name = "hlo_to_tiled_hlo", + actual = ":fusion_to_tiled_hlo", +) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tools/fusion_to_tiled_hlo.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tools/fusion_to_tiled_hlo.cc new file mode 100644 index 00000000000000..3f60be98f193aa --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tools/fusion_to_tiled_hlo.cc @@ -0,0 +1,255 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/log/check.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_split.h" +#include "absl/strings/string_view.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/MathExtras.h" +#include "llvm/Support/raw_ostream.h" +#include "mlir/IR/MLIRContext.h" +#include "xla/codegen/tiling/experimental/tiled_hlo.h" +#include "xla/codegen/tiling/experimental/tiling_space.h" +#include "xla/codegen/xtile/block_level_parameters.h" +#include "xla/codegen/xtile/tiling_from_block_parameters.h" +#include "xla/debug_options_flags.h" +#include "xla/hlo/analysis/symbolic_expr.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_traversal.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/tools/hlo_module_loader.h" +#include "xla/tsl/util/command_line_flags.h" +#include "xla/xla.pb.h" +#include "tsl/platform/init_main.h" + +namespace xla::gpu { +namespace { + +using ::xla::gpu::experimental::TiledHloComputation; +using ::xla::gpu::experimental::TilingSpace; +using ::xla::xtile::BlockLevelParameters; + +absl::StatusOr> ParseTileSizes( + absl::string_view tile_sizes_str) { + llvm::SmallVector tile_sizes; + for (absl::string_view part : + absl::StrSplit(tile_sizes_str, ',', absl::SkipEmpty())) { + int64_t size = 0; + if (!absl::SimpleAtoi(part, &size) || size <= 0) { + return absl::InvalidArgumentError( + absl::StrCat("Invalid tile size: '", part, + "'. Tile sizes must be positive integers.")); + } + tile_sizes.push_back(size); + } + return tile_sizes; +} + +absl::Status AssignTileSizesIfAvailable( + const HloInstruction& instr, + std::optional> explicit_tile_sizes, + std::unique_ptr& tiling_space) { + if (explicit_tile_sizes.has_value()) { + if (explicit_tile_sizes->size() != tiling_space->num_dimensions()) { + return absl::InvalidArgumentError(absl::StrCat( + "Number of explicit tile sizes (", explicit_tile_sizes->size(), + ") does not match number of dimensions in TilingSpace (", + tiling_space->num_dimensions(), ").")); + } + return tiling_space->AssignTileSizes(*explicit_tile_sizes); + } + + if (instr.opcode() != HloOpcode::kFusion || !instr.has_backend_config()) { + return absl::OkStatus(); + } + + auto gpu_config = instr.backend_config(); + if (!gpu_config.ok() || !gpu_config->has_fusion_backend_config() || + !gpu_config->fusion_backend_config().has_block_level_fusion_config()) { + return absl::OkStatus(); + } + + BlockLevelParameters block_level_parameters = + BlockLevelParameters::FromBlockLevelFusionConfig( + gpu_config->fusion_backend_config().block_level_fusion_config()); + const DebugOptions& debug_options = + instr.GetModule()->config().debug_options(); + bool enable_same_shape_mof = + debug_options + .xla_gpu_experimental_enable_same_shape_multi_output_fusion(); + if (block_level_parameters.output_tile_sizes.size() > 1 && + !enable_same_shape_mof && + debug_options.xla_gpu_unsupported_enable_triton_multi_output_fusion()) { + std::vector flattened; + for (const auto& sizes : block_level_parameters.output_tile_sizes) { + flattened.insert(flattened.end(), sizes.begin(), sizes.end()); + } + block_level_parameters.output_tile_sizes = {std::move(flattened)}; + } + ABSL_ASSIGN_OR_RETURN( + llvm::SmallVector tile_sizes, + xtile::GetTilingSpaceConcreteSizes(*tiling_space, block_level_parameters, + enable_same_shape_mof)); + for (int64_t& size : tile_sizes) { + size = size == 0 ? 1 : llvm::PowerOf2Ceil(size); + } + return tiling_space->AssignTileSizes(tile_sizes); +} + +absl::StatusOr TileInstruction( + const HloInstruction& instr, + std::optional> explicit_tile_sizes, + bool symbolic, mlir::MLIRContext& mlir_context) { + std::unique_ptr fusion_adaptor = + HloFusionAdaptor::ForInstruction(&instr); + ABSL_ASSIGN_OR_RETURN(std::unique_ptr tiling_space, + TilingSpace::Create(*fusion_adaptor, &mlir_context)); + if (!symbolic) { + ABSL_RETURN_IF_ERROR(AssignTileSizesIfAvailable( + instr, std::move(explicit_tile_sizes), tiling_space)); + } + ABSL_ASSIGN_OR_RETURN( + TiledHloComputation tiled_computation, + TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space))); + tiled_computation.Simplify(); + tiled_computation.SortInstructionsPostOrder(); + return tiled_computation; +} + +const HloInstruction* GetOrCreateEntryTarget(HloComputation* entry) { + HloInstruction* root = entry->root_instruction(); + bool needs_fusion_wrapper = root->opcode() == HloOpcode::kTuple; + for (const HloInstruction* operand : root->operands()) { + if (operand->opcode() != HloOpcode::kParameter && + operand->opcode() != HloOpcode::kConstant) { + needs_fusion_wrapper = true; + break; + } + } + if (!needs_fusion_wrapper) { + return root; + } + std::vector reachable = + HloBfsFindAll({root}, [](const HloInstruction*) { return true; }); + std::vector post_order = entry->MakeInstructionPostOrder(); + std::vector to_fuse; + for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) { + HloInstruction* instr = *it; + if (instr->opcode() != HloOpcode::kParameter && + absl::c_linear_search(reachable, instr)) { + to_fuse.push_back(instr); + } + } + return entry->CreateFusionInstruction(to_fuse, + HloInstruction::FusionKind::kCustom); +} + +absl::Status RealMain(absl::string_view input_file, + absl::string_view tile_sizes_str, + absl::string_view fusion_name, bool symbolic) { + ABSL_ASSIGN_OR_RETURN(std::unique_ptr hlo_module, + xla::LoadModuleFromFile(std::string(input_file))); + + std::optional> explicit_tile_sizes; + if (!tile_sizes_str.empty()) { + ABSL_ASSIGN_OR_RETURN(explicit_tile_sizes, ParseTileSizes(tile_sizes_str)); + } + + mlir::MLIRContext mlir_context; + RegisterSymbolicExprStorage(&mlir_context); + + std::vector targets; + for (const HloInstruction* instr : + hlo_module->entry_computation()->MakeInstructionPostOrder()) { + if (instr->opcode() == HloOpcode::kFusion) { + if (fusion_name.empty() || instr->name() == fusion_name || + instr->fused_instructions_computation()->name() == fusion_name) { + targets.push_back(instr); + } + } + } + + if (targets.empty()) { + if (!fusion_name.empty()) { + return absl::NotFoundError( + absl::StrCat("Fusion '", fusion_name, "' not found in module.")); + } + targets.push_back(GetOrCreateEntryTarget(hlo_module->entry_computation())); + } + + for (const HloInstruction* target : targets) { + if (targets.size() > 1) { + llvm::outs() << "Fusion: " << target->name() << "\n"; + } + ABSL_ASSIGN_OR_RETURN( + TiledHloComputation tiled_computation, + TileInstruction(*target, explicit_tile_sizes, symbolic, mlir_context)); + llvm::outs() << tiled_computation.ToString(); + } + return absl::OkStatus(); +} + +} // namespace +} // namespace xla::gpu + +int main(int argc, char** argv) { + std::string tile_sizes; + std::string fusion_name; + bool symbolic = false; + std::vector flag_list = { + tsl::Flag("tile_sizes", &tile_sizes, + "Comma-separated concrete tile sizes for TilingSpace " + "dimensions. Overrides block_level_fusion_config if set."), + tsl::Flag("fusion_name", &fusion_name, + "Optional name of the fusion instruction or fused computation " + "to tile."), + tsl::Flag("symbolic", &symbolic, + "If true, keep tile sizes symbolic even if " + "block_level_fusion_config is present."), + }; + xla::AppendDebugOptionsFlags(&flag_list); + const std::string kUsageString = tsl::Flags::Usage(argv[0], flag_list); + bool parse_ok = tsl::Flags::Parse(&argc, argv, flag_list); + tsl::port::InitMain(argv[0], &argc, &argv); + if (!parse_ok) { + // Print the usage using cerr to avoid truncation by LOG. + std::cerr << kUsageString; + return 1; + } + CHECK_GT(argc, 1) << "Must specify an input file"; + absl::Status status = + xla::gpu::RealMain(argv[1], tile_sizes, fusion_name, symbolic); + if (!status.ok()) { + // We don't return non-zero codes as some of the tests check the status. + std::cerr << status << "\n"; + } + return 0; +} diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir index 7802eea10c4025..a86502850d8b8a 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir +++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir @@ -575,3 +575,35 @@ module { // CHECK-TDM-LABEL: tt.func @apply_mask_to_aligned_offset_with_out_of_bounds_reads_at_end // CHECK-TDM: tt.descriptor_load // CHECK-TDM: tt.descriptor_store + +// ----- + +#indexing_map_reduced_oob = #xla.indexing_map<"(pid) -> (pid floordiv 2), domain: pid in [0, 7]"> +module { + func.func @apply_mask_to_reduced_dim_with_out_of_bounds_offset(%arg0: !tt.ptr, %arg1: !tt.ptr) { + %0 = tt.get_program_id x : i32 + %1 = arith.index_cast %0 : i32 to index + %2 = xla.apply_indexing #indexing_map_reduced_oob(%1) + // Dimension 0 has size 3, while %2 ranges in [0, 3]. Even though the tile + // size along dimension 0 is 1 and rank-reduced, a bounds mask must still be + // emitted for dimension 0. + %extracted_tile = triton_xla.extract from %arg0 + as memref<3x8xbf16, #xtile.layout<[1, 0]>> + [%2, 0] [1, 8] [1, 1] : tensor<8xbf16> + triton_xla.insert %extracted_tile into %arg1 + as memref<4x8xbf16, #xtile.layout<[1, 0]>> + [%2, 0] [1, 8] [1, 1] : tensor<8xbf16> + func.return + } +} + +// CHECK-LABEL: tt.func @apply_mask_to_reduced_dim_with_out_of_bounds_offset +// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : i64 +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i64 +// CHECK: %[[OFFSET:.*]] = arith.index_cast %{{.*}} : index to i64 +// CHECK: %[[RIGHT_MASK:.*]] = arith.cmpi slt, %[[OFFSET]], %[[C3]] : i64 +// CHECK: %[[LEFT_MASK:.*]] = arith.cmpi sge, %[[OFFSET]], %[[C0]] : i64 +// CHECK: %[[MASK:.*]] = arith.andi %[[LEFT_MASK]], %[[RIGHT_MASK]] : i1 +// CHECK: %[[SPLAT_MASK:.*]] = tt.splat %[[MASK]] : i1 -> tensor<8xi1> +// CHECK: tt.load {{.*}}, %[[SPLAT_MASK]], {{.*}} +// CHECK: tt.store {{.*}}, %{{.*}} : tensor<8x!tt.ptr> diff --git a/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h b/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h index 81df37f3127b0d..f999d82bf71344 100644 --- a/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h +++ b/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h @@ -240,6 +240,18 @@ class GpuCommunicator : public Communicator { RankId root, const Executor& executor) = 0; + // Reduces buffers of length `count` in `send_buffer` using `reduction_kind` + // and writes the result only to `recv_buffer` on the `root` rank (NCCL + // `ncclReduce` semantics). On non-root ranks `recv_buffer` is left untouched. + // Defaults to unimplemented so backends can opt in; NCCL overrides it. + virtual absl::Status LaunchReduce(se::DeviceAddressBase send_buffer, + se::DeviceAddressBase recv_buffer, + PrimitiveType dtype, size_t count, + ReductionKind reduction_kind, RankId root, + const Executor& executor) { + return Unimplemented("LaunchReduce is not implemented"); + } + virtual absl::Status LaunchReduceScatter(se::DeviceAddressBase send_buffer, se::DeviceAddressBase recv_buffer, PrimitiveType dtype, size_t count, diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc index cad31b7fe963e4..ba521525d76a6a 100644 --- a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc +++ b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc @@ -626,6 +626,45 @@ absl::Status NcclCommunicator::LaunchBroadcast( return absl::OkStatus(); } +absl::Status NcclCommunicator::LaunchReduce(se::DeviceAddressBase send_buffer, + se::DeviceAddressBase recv_buffer, + PrimitiveType dtype, size_t count, + ReductionKind reduction_kind, + RankId root, + const Executor& executor) { + if (cancel_->IsCancelled()) { + return FailedPrecondition("NcclCommunicator aborted"); + } + se::Stream* stream = ToStream(executor); + + { + absl::MutexLock lock(comm_->mutex); + XLA_VLOG_DEVICE(3, stream->parent()->device_ordinal()) + << absl::StreamFormat( + "Launch NCCL Reduce operation; send_buffer=%p; " + "recv_buffer=%p; dtype=%s; count=%d; reduction_kind=%v; " + "root=%d; comm=%p; stream=%p", + send_buffer.opaque(), recv_buffer.opaque(), + primitive_util::LowercasePrimitiveTypeName(dtype), count, + reduction_kind, root.value(), comm_->comm, stream); + + ABSL_ASSIGN_OR_RETURN(ncclDataType_t nccl_dtype, + ToNcclDataType(dtype, /*is_reduction_op=*/true, + stream->parent() + ->GetDeviceDescription() + .cuda_compute_capability())); + + ABSL_RETURN_IF_ERROR(XLA_NCCL_STATUS(ncclReduce( + send_buffer.opaque(), recv_buffer.opaque(), ToNcclCount(dtype, count), + nccl_dtype, ToNcclReduction(reduction_kind), root.value(), comm_->comm, + AsCudaStream(stream)))); + } + if (!IsInsideNcclGroupLaunch()) { + ABSL_RETURN_IF_ERROR(PollUntilDone()); + } + return absl::OkStatus(); +} + absl::Status NcclCommunicator::LaunchReduceScatter( se::DeviceAddressBase send_buffer, se::DeviceAddressBase recv_buffer, PrimitiveType dtype, size_t count, ReductionKind reduction_kind, diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h index b6f132d9c7519f..ef343f18be5351 100644 --- a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h +++ b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h @@ -217,6 +217,12 @@ class NcclCommunicator : public GpuCommunicator { PrimitiveType dtype, size_t count, RankId root, const Executor& executor) final; + absl::Status LaunchReduce(se::DeviceAddressBase send_buffer, + se::DeviceAddressBase recv_buffer, + PrimitiveType dtype, size_t count, + ReductionKind reduction_kind, RankId root, + const Executor& executor) final; + absl::Status LaunchReduceScatter(se::DeviceAddressBase send_buffer, se::DeviceAddressBase recv_buffer, PrimitiveType dtype, size_t count, diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD index 7a570f6651e4da..98ef2398d0bf7e 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/BUILD @@ -19,6 +19,7 @@ load("@rules_cc//cc:cc_library.bzl", "cc_library") load("//xla:xla.default.bzl", "xla_cc_test") +load("//xla/tsl:tsl.bzl", "internal_visibility") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], @@ -45,10 +46,10 @@ package_group( cc_library( name = "native_custom_call_emitter_context", hdrs = ["native_custom_call_emitter_context.h"], - visibility = [ + visibility = internal_visibility([ ":handler_allowlist", "//xla/service/gpu:__pkg__", - ], + ]), deps = [ "//xla:shape_util", "//xla:xla_proto_cc", @@ -87,7 +88,7 @@ cc_library( name = "native_custom_call_handler_registration", srcs = ["native_custom_call_handler_registration.cc"], hdrs = ["native_custom_call_handler_registration.h"], - visibility = [":handler_allowlist"], + visibility = internal_visibility([":handler_allowlist"]), deps = [ ":native_custom_call_handler_registry", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 31122dd1cba818..0d3efec73c29e9 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -137,7 +137,6 @@ cc_library( "//xla/stream_executor:command_buffer", "//xla/stream_executor:platform", "@com_google_absl//absl/functional:function_ref", - "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", @@ -351,10 +350,12 @@ xla_cc_test( name = "command_buffer_cmd_emitter_test", srcs = ["command_buffer_cmd_emitter_test.cc"], deps = [ + ":async_thunk", ":command", ":command_buffer_cmd_emitter", ":command_executor", ":conditional_thunk", + ":execution_stream_id", ":kernel_thunk", ":sequential_thunk", ":thunk", @@ -366,13 +367,16 @@ xla_cc_test( "//xla/runtime:execution_graph", "//xla/service:buffer_assignment", "//xla/service:shaped_slice", - "//xla/tsl/platform:statusor", + "//xla/service/gpu:launch_dimensions", + "//xla/stream_executor:launch_dim", + "//xla/stream_executor/gpu:tma_metadata", "//xla/tsl/platform:test", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", + "@tsl//tsl/platform:status_matchers", ], ) @@ -2063,6 +2067,46 @@ xla_test( ], ) +cc_library( + name = "collective_reduce_thunk", + srcs = ["collective_reduce_thunk.cc"], + hdrs = ["collective_reduce_thunk.h"], + deps = [ + ":all_reduce_thunk", + ":collective_thunk", + ":collective_thunk_proto_cc", + ":thunk", + ":thunk_proto_cc", + "//xla:future", + "//xla:util", + "//xla:xla_data_proto_cc", + "//xla/backends/gpu/collectives:gpu_clique_key", + "//xla/backends/gpu/collectives:gpu_collectives", + "//xla/backends/gpu/collectives:gpu_communicator", + "//xla/core/collectives:communicator", + "//xla/core/collectives:rank_id", + "//xla/core/collectives:reduction_kind", + "//xla/hlo/ir:hlo", + "//xla/service:buffer_assignment", + "//xla/service:collective_ops_utils", + "//xla/stream_executor:device_address", + "//xla/stream_executor:memory_allocation", + "//xla/stream_executor:stream", + "//xla/stream_executor:stream_executor_h", + "//xla/tsl/platform:logging", + "@com_google_absl//absl/base", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:span", + ], +) + cc_library( name = "all_reduce_thunk", srcs = ["all_reduce_thunk.cc"], @@ -2594,6 +2638,36 @@ xla_test( ], ) +xla_test( + name = "collective_reduce_thunk_multigpu_test", + srcs = ["collective_reduce_thunk_multigpu_test.cc"], + backend_tags = { + "gpu": ["multi_gpu"], + }, + backends = ["gpu"], + deps = [ + ":all_reduce_thunk", + ":collective_reduce_thunk", + ":collective_thunk", + ":collective_thunk_multigpu_test_utils", + ":thunk", + "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla/core/collectives:reduction_kind", + "//xla/service:buffer_assignment", + "//xla/service:device_assignment", + "//xla/service:shaped_slice", + "//xla/service/gpu:buffer_allocations", + "//xla/stream_executor:device_address", + "//xla/stream_executor:stream", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest_main", + ], +) + xla_test( name = "collective_broadcast_thunk_multigpu_test", srcs = ["collective_broadcast_thunk_multigpu_test.cc"], @@ -2895,13 +2969,11 @@ cc_library( "//xla/stream_executor:stream", "//xla/stream_executor:stream_executor_h", "//xla/tsl/platform:statusor", - "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:casts", ], @@ -4245,6 +4317,7 @@ xla_cc_test( "//xla/tsl/lib/gtl:int_type", "//xla/tsl/platform:test", "//xla/tsl/platform:test_main", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", @@ -4544,6 +4617,7 @@ cc_library( ":collective_group_thunk", ":collective_kernel_thunk", ":collective_permute_thunk", + ":collective_reduce_thunk", ":conditional_thunk", ":convolution_reorder_thunk", ":convolution_thunk", diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk_multigpu_test.cc index e9c35218d362c8..2f0676ef6b1f78 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk_multigpu_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk_multigpu_test.cc @@ -43,9 +43,7 @@ namespace xla::gpu { namespace { static constexpr int kNumDevices = 2; -static constexpr int64_t kLength = 4; -static constexpr int64_t kByteLength = sizeof(float) * kLength; -static constexpr int64_t kGatheredLength = kLength * kNumDevices; +static constexpr int64_t kGatheredLength = kNumElements * kNumDevices; static constexpr int64_t kGatheredByteLength = sizeof(float) * kGatheredLength; static CollectiveConfig MakeAllGatherConfig() { @@ -63,12 +61,13 @@ static CollectiveConfig MakeAllGatherConfig() { static AllGatherThunk MakeThunk(const BufferAllocation& alloc_src, const BufferAllocation& alloc_dst) { - ShapedSlice src_slice{BufferAllocation::Slice(&alloc_src, 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; + ShapedSlice src_slice{ + BufferAllocation::Slice(&alloc_src, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; ShapedSlice dst_slice{ BufferAllocation::Slice(&alloc_dst, 0, kGatheredByteLength), ShapeUtil::MakeShape(F32, {kGatheredLength})}; - CollectiveThunk::Buffer buffer{.element_count = kLength, + CollectiveThunk::Buffer buffer{.element_count = kNumElements, .source_buffer = src_slice, .destination_buffer = dst_slice, .source_memory_space = 0, @@ -79,12 +78,12 @@ static AllGatherThunk MakeThunk(const BufferAllocation& alloc_src, using DeviceTestSlot = CollectiveThunkMultiGpuTestState; static std::vector DeviceBufferSizes() { - return {kByteLength, kGatheredByteLength}; + return {kFloatByteLength, kGatheredByteLength}; } static std::vector SourceValues(int source_rank, int phase) { - std::vector values(kLength); - for (int i = 0; i < kLength; ++i) { + std::vector values(kNumElements); + for (int i = 0; i < kNumElements; ++i) { values[i] = static_cast(phase * 100 + source_rank * 10 + i); } return values; @@ -116,8 +115,8 @@ static absl::Status VerifyOutput(se::Stream& stream, se::DeviceAddressBase dst, ReadDeviceBuffer(stream, dst, kGatheredLength)); for (int source_rank = 0; source_rank < kNumDevices; ++source_rank) { std::vector expected = SourceValues(source_rank, phase); - for (int i = 0; i < kLength; ++i) { - int output_index = source_rank * kLength + i; + for (int i = 0; i < kNumElements; ++i) { + int output_index = source_rank * kNumElements + i; if (output[output_index] != expected[i]) { return absl::InternalError( absl::StrFormat("output[%d] = %g, expected %g", output_index, @@ -191,7 +190,7 @@ TEST(AllGatherThunkMultiGpuTest, ExecuteOnStream) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); BufferAllocation alloc_dst(/*index=*/1, kGatheredByteLength, /*color=*/0); AllGatherThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -212,7 +211,7 @@ TEST(AllGatherThunkMultiGpuTest, RecordCommandBufferCreate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); BufferAllocation alloc_dst(/*index=*/1, kGatheredByteLength, /*color=*/0); AllGatherThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -233,7 +232,7 @@ TEST(AllGatherThunkMultiGpuTest, RecordCommandBufferUpdate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); BufferAllocation alloc_dst(/*index=*/1, kGatheredByteLength, /*color=*/0); AllGatherThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc index 2620c32d256c6a..1e46607f1d5216 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc @@ -20,11 +20,13 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/status_matchers.h" +#include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" @@ -60,6 +62,18 @@ using ::testing::HasSubstr; TSL_LIB_GTL_DEFINE_INT_TYPE(CollectiveKernelEnabled, bool); TSL_LIB_GTL_DEFINE_INT_TYPE(MultimemEnabled, bool); +// Number of devices per host in the test topology. Large enough for all replica +// groups used in these tests to be local, including groups that exceed the +// maximum number of ranks supported by the all-reduce kernel. +constexpr int32_t kNumDevicesPerHost = 2 * se::gpu::kMaxNumAllReduceInputPtrs; + +// Returns a replica group containing replicas [0, num_replicas). +std::vector IotaReplicaGroup(int64_t num_replicas) { + std::vector replica_group(num_replicas); + absl::c_iota(replica_group, 0); + return replica_group; +} + class BuildAllReduceInfoTest : public HloHardwareIndependentTestBase { protected: // Helper to reduce boilerplate while keeping tests independent. Supports @@ -93,7 +107,8 @@ class BuildAllReduceInfoTest : public HloHardwareIndependentTestBase { gpu::GpuTargetConfig::FromProto(target_config_proto)); GpuTopology gpu_topology("platform_version", /*num_partitions=*/1, /*num_hosts_per_partition=*/1, - /*num_devices_per_host=*/16, target_config); + /*num_devices_per_host=*/kNumDevicesPerHost, + target_config); int64_t num_replicas = 0; std::vector group_strs; group_strs.reserve(replica_groups.size()); @@ -177,12 +192,23 @@ TEST_F(BuildAllReduceInfoTest, FailsForNonPowerOfTwoDevices) { HasSubstr("only supported for power of 2"))); } -TEST_F(BuildAllReduceInfoTest, FailsForTooManyDevices) { +TEST_F(BuildAllReduceInfoTest, SupportsMaxNumDevices) { EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false), F32, {1024}, HloOpcode::kAdd, - {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}), - StatusIs(absl::StatusCode::kUnimplemented, - HasSubstr("does not support more than 8 ranks"))); + IotaReplicaGroup(se::gpu::kMaxNumAllReduceInputPtrs)), + IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy, + AllReduceStrategy::kOneShot))); +} + +TEST_F(BuildAllReduceInfoTest, FailsForTooManyDevices) { + EXPECT_THAT( + BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false), F32, + {1024}, HloOpcode::kAdd, + IotaReplicaGroup(2 * se::gpu::kMaxNumAllReduceInputPtrs)), + StatusIs(absl::StatusCode::kUnimplemented, + HasSubstr(absl::StrCat("does not support more than ", + se::gpu::kMaxNumAllReduceInputPtrs, + " ranks")))); } TEST_F(BuildAllReduceInfoTest, FailsForUnsupportedTypeCombination) { diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce_thunk_multigpu_test.cc index 6255e4203dd9d6..f5d9aeaa1d9617 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_reduce_thunk_multigpu_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce_thunk_multigpu_test.cc @@ -45,10 +45,8 @@ namespace xla::gpu { namespace { static constexpr int kNumDevices = 2; -static constexpr int64_t kLength = 4; -static constexpr int64_t kByteLength = sizeof(float) * kLength; -static_assert(kLength % kNumDevices == 0); -static constexpr int64_t kReduceScatterLength = kLength / kNumDevices; +static_assert(kNumElements % kNumDevices == 0); +static constexpr int64_t kReduceScatterLength = kNumElements / kNumDevices; static constexpr int64_t kReduceScatterByteLength = sizeof(float) * kReduceScatterLength; @@ -115,14 +113,14 @@ static DirectAllReduceThunk MakeAllReduceThunk( const BufferAllocation& alloc_src, const BufferAllocation& alloc_dst) { return DirectAllReduceThunk( Thunk::ThunkInfo(), MakeSumConfig(), - {MakeBuffer(alloc_src, alloc_dst, kLength, kLength)}); + {MakeBuffer(alloc_src, alloc_dst, kNumElements, kNumElements)}); } static ReduceScatterThunk MakeReduceScatterThunk( const BufferAllocation& alloc_src, const BufferAllocation& alloc_dst) { return ReduceScatterThunk( Thunk::ThunkInfo(), MakeSumConfig(), - {MakeBuffer(alloc_src, alloc_dst, kLength, kReduceScatterLength)}); + {MakeBuffer(alloc_src, alloc_dst, kNumElements, kReduceScatterLength)}); } using DeviceTestSlot = CollectiveThunkMultiGpuTestState; @@ -151,7 +149,7 @@ static float DeviceScaleSum() { static std::vector AllReduceInput(int device_ordinal, float phase_scale) { return std::vector( - kLength, static_cast(device_ordinal + 1) * phase_scale); + kNumElements, static_cast(device_ordinal + 1) * phase_scale); } static std::vector AllReduceExpected(float phase_scale) { @@ -159,15 +157,15 @@ static std::vector AllReduceExpected(float phase_scale) { for (int d = 0; d < kNumDevices; ++d) { sum += static_cast(d + 1) * phase_scale; } - return std::vector(kLength, sum); + return std::vector(kNumElements, sum); } static std::vector ReduceScatterInput(int device_ordinal, float phase_scale) { std::vector data; - data.reserve(kLength); + data.reserve(kNumElements); float scale = DeviceScale(device_ordinal) * phase_scale; - for (int i = 0; i < kLength; ++i) { + for (int i = 0; i < kNumElements; ++i) { data.push_back(static_cast(i + 1) * scale); } return data; @@ -211,7 +209,7 @@ static std::vector ExpectedValues(CollectiveTestKind kind, static int64_t DestinationByteLength(CollectiveTestKind kind) { switch (kind) { case CollectiveTestKind::kAllReduce: - return kByteLength; + return kFloatByteLength; case CollectiveTestKind::kReduceScatter: return kReduceScatterByteLength; } @@ -225,7 +223,7 @@ static absl::Status SetupDeviceSlot(int device_ordinal, DeviceTestSlot& slot, AllReduceReduceScatterThunkBase& thunk, const DeviceAssignment& device_assignment, CollectiveTestKind kind) { - std::vector buffer_sizes = {kByteLength, + std::vector buffer_sizes = {kFloatByteLength, DestinationByteLength(kind)}; return SetupCollectiveThunkDevice(device_ordinal, kNumDevices, buffer_sizes, thunk, device_assignment, slot); @@ -320,8 +318,8 @@ TEST(AllReduceThunkMultiGpuTest, RecordCommandBufferCreate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); DirectAllReduceThunk thunk = MakeAllReduceThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -345,8 +343,8 @@ TEST(AllReduceThunkMultiGpuTest, RecordCommandBufferUpdate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); DirectAllReduceThunk thunk = MakeAllReduceThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -381,7 +379,7 @@ TEST(ReduceScatterThunkMultiGpuTest, RecordCommandBufferCreate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); BufferAllocation alloc_dst(/*index=*/1, kReduceScatterByteLength, /*color=*/0); ReduceScatterThunk thunk = MakeReduceScatterThunk(alloc_src, alloc_dst); @@ -407,7 +405,7 @@ TEST(ReduceScatterThunkMultiGpuTest, RecordCommandBufferUpdate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); BufferAllocation alloc_dst(/*index=*/1, kReduceScatterByteLength, /*color=*/0); ReduceScatterThunk thunk = MakeReduceScatterThunk(alloc_src, alloc_dst); diff --git a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc index edea37b614a4d8..1b4babf5f995dd 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc @@ -44,8 +44,6 @@ namespace xla::gpu { namespace { static constexpr int kNumDevices = 2; -static constexpr int64_t kLength = 4; -static constexpr int64_t kByteLength = sizeof(float) * kLength; static constexpr int kNumBuffers = kNumDevices; static AllToAllConfig MakeAllToAllConfig() { @@ -67,7 +65,7 @@ static std::vector MakeThunkBufferAllocations() { std::vector allocations; allocations.reserve(2 * kNumBuffers); for (int i = 0; i < 2 * kNumBuffers; ++i) { - allocations.emplace_back(/*index=*/i, kByteLength, /*color=*/0); + allocations.emplace_back(/*index=*/i, kFloatByteLength, /*color=*/0); } return allocations; } @@ -77,12 +75,12 @@ static AllToAllThunk MakeThunk(absl::Span allocations) { buffers.reserve(kNumBuffers); for (int i = 0; i < kNumBuffers; ++i) { ShapedSlice src_slice{ - BufferAllocation::Slice(&allocations[i], 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - ShapedSlice dst_slice{ - BufferAllocation::Slice(&allocations[kNumBuffers + i], 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - buffers.push_back(CollectiveThunk::Buffer{.element_count = kLength, + BufferAllocation::Slice(&allocations[i], 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + ShapedSlice dst_slice{BufferAllocation::Slice(&allocations[kNumBuffers + i], + 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + buffers.push_back(CollectiveThunk::Buffer{.element_count = kNumElements, .source_buffer = src_slice, .destination_buffer = dst_slice, .source_memory_space = 0, @@ -97,7 +95,7 @@ static AllToAllThunk MakeThunk(absl::Span allocations) { using DeviceTestSlot = CollectiveThunkMultiGpuTestState; static std::vector DeviceBufferSizes() { - return std::vector(2 * kNumBuffers, kByteLength); + return std::vector(2 * kNumBuffers, kFloatByteLength); } static std::vector AllAllocationIndices() { @@ -126,7 +124,7 @@ static float SourceValue(int source_rank, int target_rank, int phase) { static absl::Status FillDeviceBufferWithValue(se::Stream& stream, se::DeviceAddressBase buffer, float value) { - std::vector data(kLength, value); + std::vector data(kNumElements, value); return FillDeviceBuffer(stream, buffer, data); } @@ -165,8 +163,8 @@ static absl::Status VerifyOutput(se::Stream& stream, float expected = SourceValue(source_rank, /*target_rank=*/device_ordinal, phase); ABSL_ASSIGN_OR_RETURN(std::vector output, - ReadDeviceBuffer(stream, dst[source_rank], kLength)); - for (int i = 0; i < kLength; ++i) { + ReadDeviceBuffer(stream, dst[source_rank], kNumElements)); + for (int i = 0; i < kNumElements; ++i) { if (output[i] != expected) { return absl::InternalError(absl::StrFormat( "dst[%d][%d] on device %d = %g, expected %g", source_rank, i, diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_multigpu_test.cc index 36110d32fb1859..0072add6bd5875 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_multigpu_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_multigpu_test.cc @@ -44,8 +44,6 @@ namespace xla::gpu { namespace { static constexpr int kNumDevices = 2; -static constexpr int64_t kLength = 4; -static constexpr int64_t kByteLength = sizeof(float) * kLength; static CollectiveConfig MakeCollectiveBroadcastConfig( bool has_dynamic_root = false) { @@ -67,11 +65,13 @@ static CollectiveConfig MakeCollectiveBroadcastConfig( static CollectiveBroadcastThunk MakeThunk( const BufferAllocation& alloc_src, const BufferAllocation& alloc_dst, const BufferAllocation* alloc_root = nullptr) { - ShapedSlice src_slice{BufferAllocation::Slice(&alloc_src, 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - ShapedSlice dst_slice{BufferAllocation::Slice(&alloc_dst, 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - CollectiveThunk::Buffer buffer{.element_count = kLength, + ShapedSlice src_slice{ + BufferAllocation::Slice(&alloc_src, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + ShapedSlice dst_slice{ + BufferAllocation::Slice(&alloc_dst, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + CollectiveThunk::Buffer buffer{.element_count = kNumElements, .source_buffer = src_slice, .destination_buffer = dst_slice, .source_memory_space = 0, @@ -97,12 +97,12 @@ static CollectiveBroadcastThunk MakeThunk( using DeviceTestSlot = CollectiveThunkMultiGpuTestState; static std::vector DeviceBufferSizes() { - return {kByteLength, kByteLength}; + return {kFloatByteLength, kFloatByteLength}; } static std::vector SourceValues(int device_ordinal, int phase) { - std::vector values(kLength); - for (int i = 0; i < kLength; ++i) { + std::vector values(kNumElements); + for (int i = 0; i < kNumElements; ++i) { values[i] = static_cast(phase * 100 + device_ordinal * 10 + i); } return values; @@ -117,7 +117,8 @@ static absl::Status FillSourceBuffer(se::Stream& stream, static absl::Status FillDestinationBuffer(se::Stream& stream, se::DeviceAddressBase buffer, float value) { - return FillDeviceBuffer(stream, buffer, std::vector(kLength, value)); + return FillDeviceBuffer(stream, buffer, + std::vector(kNumElements, value)); } static absl::Status PrepareInputs( @@ -130,9 +131,9 @@ static absl::Status PrepareInputs( static absl::Status VerifyOutput(se::Stream& stream, se::DeviceAddressBase dst, int device_ordinal, int phase) { ABSL_ASSIGN_OR_RETURN(std::vector output, - ReadDeviceBuffer(stream, dst, kLength)); + ReadDeviceBuffer(stream, dst, kNumElements)); std::vector expected = SourceValues(/*device_ordinal=*/0, phase); - for (int i = 0; i < kLength; ++i) { + for (int i = 0; i < kNumElements; ++i) { if (output[i] != expected[i]) { return absl::InternalError( absl::StrFormat("device %d output[%d] = %g, expected %g", @@ -214,8 +215,8 @@ TEST(CollectiveBroadcastThunkMultiGpuTest, ExecuteOnStream) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); CollectiveBroadcastThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -235,8 +236,8 @@ TEST(CollectiveBroadcastThunkMultiGpuTest, RecordCommandBufferCreate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); CollectiveBroadcastThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -256,8 +257,8 @@ TEST(CollectiveBroadcastThunkMultiGpuTest, RecordCommandBufferUpdate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); CollectiveBroadcastThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -278,8 +279,8 @@ TEST(CollectiveBroadcastThunkMultiGpuTest, ExecuteOnStreamWithDynamicRoot) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); BufferAllocation alloc_root(/*index=*/2, sizeof(int32_t), /*color=*/0); CollectiveBroadcastThunk thunk = MakeThunk(alloc_src, alloc_dst, &alloc_root); diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc index cd61489fdfa30f..223d7479a5f77a 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc @@ -651,7 +651,10 @@ absl::Status CollectiveKernelThunk::ExecuteOnStream( state = it->second.get(); } - state->invocation_count += kernel_spec_.codegen_config.sync_count_increment; + // Backward-compatibility fallback for legacy AOT-compiled kernels. + if (!kernel_spec_.codegen_config.device_sync_count) { + state->invocation_count += kernel_spec_.codegen_config.sync_count_increment; + } TF_RET_CHECK(state->kernel != nullptr) << "Kernel is not initialized for collective kernel thunk."; @@ -781,6 +784,8 @@ CollectiveKernelThunk::FromProto( proto_spec.invocation_count_increment(); kernel_spec.codegen_config.copy_input_to_scratch = proto_spec.copy_input_to_scratch(); + kernel_spec.codegen_config.device_sync_count = + proto_spec.device_sync_count(); } else { // Backward-compatibility fallback for legacy AOT-compiled kernels without // an explicit CollectiveKernelSpec. @@ -828,6 +833,8 @@ CollectiveKernelThunk::FromProto( kernel_spec.codegen_config.sync_count_increment = 1 + static_cast(GetAllReduceStrategy( input_size_bytes, /*is_multimem_enabled=*/false)); + // Legacy kernels read the host-provided invocation count. + kernel_spec.codegen_config.device_sync_count = false; } std::vector buffers; buffers.reserve(thunk_proto.buffers_size()); @@ -900,6 +907,8 @@ absl::StatusOr CollectiveKernelThunk::ToProto() const { kernel_spec_.codegen_config.sync_count_increment); proto_spec->set_copy_input_to_scratch( kernel_spec_.codegen_config.copy_input_to_scratch); + proto_spec->set_device_sync_count( + kernel_spec_.codegen_config.device_sync_count); for (const CollectiveThunk::Buffer& buffer : buffers_) { ABSL_ASSIGN_OR_RETURN(*thunk_proto->add_buffers(), buffer.ToProto()); diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto index 3d1ae344248142..16288728c4c76d 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto @@ -63,6 +63,12 @@ message CollectiveKernelSpecProto { repeated KernelArgDescriptorProto argument_descriptors = 4; int32 invocation_count_increment = 5; bool copy_input_to_scratch = 6; + // If true, the kernel derives its invocation counter from its own barrier + // slot in device memory (`SignalBuffers[rank][block_id * world_size + rank]`) + // and ignores the invocation count argument. Unset (false) for legacy kernels + // that read the invocation count argument. + // Will be deprecated and enabled by default by 01/03/2026. + bool device_sync_count = 7; } message CollectiveKernelThunkProto { diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc index 5f16b3a0899220..79192758a8ddb9 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc @@ -83,24 +83,28 @@ using ::absl_testing::StatusIs; using ::testing::ElementsAre; static constexpr absl::string_view kProfileName = "test_kernel_profiler"; -static constexpr absl::string_view kKernelName = "six_argument_kernel"; +static constexpr absl::string_view kKernelName = "five_argument_kernel"; static constexpr int64_t kNumElements = 128; // Test kernel was compiled using following CUDA source: -// __global__ void six_argument_kernel(int64_t* input_buffer, // 1 -// int64_t* output_buffer, // 2 -// int64_t rank, // 3 -// int64_t signal_value // 4 -// int64_t* signal_buffers, // 5 -// int64_t* remote_buffers, // 6 +// __global__ void five_argument_kernel(int64_t* input_buffer, // 1 +// int64_t* output_buffer, // 2 +// int64_t rank, // 3 +// int32_t** signal_buffers, // 4 +// int64_t* remote_buffers // 5 // ) { -// (void)rank; -// (void)signal_buffers; // (void)remote_buffers; +// // Mirrors EmitDeviceInvocationCount: the invocation count is derived from +// // this block's barrier slot SignalBuffers[rank][blockIdx.x * world_size + +// // rank], without a thread barrier or a store. The test launches a single +// // block, so the slot index reduces to `rank`. +// int32_t* local_signal_buffer = signal_buffers[rank]; +// int32_t signal_count = +// *(volatile int32_t*)&local_signal_buffer[rank] + 1; // int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; // for (int i = idx; i < kNumElements; i += gridDim.x * blockDim.x) { // if (i < kNumElements) { -// output_buffer[i] = input_buffer[i] + signal_value; +// output_buffer[i] = input_buffer[i] + signal_count; // } // } // } @@ -109,23 +113,34 @@ static constexpr absl::string_view kKernelSource = R"( .target sm_90 .address_size 64 - .visible .entry six_argument_kernel( + .visible .entry five_argument_kernel( .param .u64 .ptr .align 1 input_buffer, .param .u64 .ptr .align 1 output_buffer, .param .u64 rank, - .param .u64 signal_value, .param .u64 .ptr .align 1 signal_buffers, .param .u64 .ptr .align 1 remote_buffers ) { .reg .pred %p<3>; - .reg .b32 %r<7>; - .reg .b64 %rd<11>; + .reg .b32 %r<8>; + .reg .b64 %rd<14>; ld.param.b64 %rd4, [input_buffer]; cvta.to.global.u64 %rd1, %rd4; ld.param.b64 %rd5, [output_buffer]; cvta.to.global.u64 %rd2, %rd5; + ld.param.u32 %r7, [rank]; + ld.param.b64 %rd11, [signal_buffers]; + cvta.to.global.u64 %rd11, %rd11; + mul.wide.u32 %rd12, %r7, 8; + add.s64 %rd11, %rd11, %rd12; + ld.global.u64 %rd13, [%rd11]; + cvta.to.global.u64 %rd13, %rd13; + mul.wide.u32 %rd12, %r7, 4; + add.s64 %rd13, %rd13, %rd12; + ld.volatile.global.s32 %r7, [%rd13]; + add.s32 %r7, %r7, 1; + cvt.s64.s32 %rd3, %r7; mov.u32 %r3, %ctaid.x; mov.u32 %r1, %ntid.x; mov.u32 %r4, %tid.x; @@ -133,7 +148,6 @@ static constexpr absl::string_view kKernelSource = R"( setp.gt.s32 %p1, %r6, 127; @%p1 bra $L__BB0_3; // - ld.param.b64 %rd3, [signal_value]; mov.u32 %r5, %nctaid.x; mul.lo.s32 %r2, %r5, %r1; $L__BB0_2: // @@ -198,7 +212,6 @@ CollectiveKernelSpec CreateCollectiveKernelSpec( {{KernelArgType::kInputBuffer, 0}, {KernelArgType::kOutputBuffer, 0}, {KernelArgType::kRuntimeRank}, - {KernelArgType::kInvocationCount}, {KernelArgType::kScratchBuffer, 0}, {KernelArgType::kScratchBuffer, 1}}, }, diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_params.h b/third_party/xla/xla/backends/gpu/runtime/collective_params.h index b6d21c4d5624a1..2c2e97224e1f13 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_params.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_params.h @@ -178,6 +178,17 @@ struct CollectiveCodegenConfig { // this amount. For one-shot collectives, this is 1. For two-shot collectives, // this is 2 and so on. uint32_t sync_count_increment = 1; + // If true, the kernel derives its barrier signal value from device memory and + // ignores the host-provided invocation count argument, so the runtime does + // not advance the host-side count. Each block reads the last signal value it + // posted from its own barrier slot (`SignalBuffers[rank][block_id * + // world_size + rank]`) in the local rank's signal buffer (scratch buffer 0), + // which `BlockBarrierOp` updates on every barrier synchronization. The signal + // buffer is zeroed only once, when the thunk is initialized, so the counter + // persists across executions and kernels can be replayed from CUDA graphs + // (including device-side loops) without updating their arguments. False only + // for legacy kernels that read the invocation count argument. + bool device_sync_count = true; }; // This structure contains the information required to configure and launch a diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_multigpu_test.cc index 4f14c6c48086eb..6cfe721a3152a8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_multigpu_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_multigpu_test.cc @@ -45,8 +45,6 @@ namespace xla::gpu { namespace { static constexpr int kNumDevices = 2; -static constexpr int64_t kLength = 4; -static constexpr int64_t kByteLength = sizeof(float) * kLength; static P2PConfig MakeCollectivePermuteConfig() { P2PConfig config; @@ -69,11 +67,13 @@ static P2PConfig MakeCollectivePermuteConfig() { static CollectivePermuteThunk MakeThunk(const BufferAllocation& alloc_src, const BufferAllocation& alloc_dst) { - ShapedSlice src_slice{BufferAllocation::Slice(&alloc_src, 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - ShapedSlice dst_slice{BufferAllocation::Slice(&alloc_dst, 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - CollectiveThunk::Buffer buffer{.element_count = kLength, + ShapedSlice src_slice{ + BufferAllocation::Slice(&alloc_src, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + ShapedSlice dst_slice{ + BufferAllocation::Slice(&alloc_dst, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + CollectiveThunk::Buffer buffer{.element_count = kNumElements, .source_buffer = src_slice, .destination_buffer = dst_slice, .source_memory_space = 0, @@ -87,12 +87,12 @@ static CollectivePermuteThunk MakeThunk(const BufferAllocation& alloc_src, using DeviceTestSlot = CollectiveThunkMultiGpuTestState; static std::vector DeviceBufferSizes() { - return {kByteLength, kByteLength}; + return {kFloatByteLength, kFloatByteLength}; } static std::vector SourceValues(int source_rank, int phase) { - std::vector values(kLength); - for (int i = 0; i < kLength; ++i) { + std::vector values(kNumElements); + for (int i = 0; i < kNumElements; ++i) { values[i] = static_cast(phase * 100 + source_rank * 10 + i); } return values; @@ -107,7 +107,8 @@ static absl::Status FillSourceBuffer(se::Stream& stream, static absl::Status FillDestinationBuffer(se::Stream& stream, se::DeviceAddressBase buffer, float value) { - return FillDeviceBuffer(stream, buffer, std::vector(kLength, value)); + return FillDeviceBuffer(stream, buffer, + std::vector(kNumElements, value)); } static absl::Status PrepareInputs( @@ -120,10 +121,10 @@ static absl::Status PrepareInputs( static absl::Status VerifyOutput(se::Stream& stream, se::DeviceAddressBase dst, int device_ordinal, int phase) { ABSL_ASSIGN_OR_RETURN(std::vector output, - ReadDeviceBuffer(stream, dst, kLength)); + ReadDeviceBuffer(stream, dst, kNumElements)); int expected_source_rank = 1 - device_ordinal; std::vector expected = SourceValues(expected_source_rank, phase); - for (int i = 0; i < kLength; ++i) { + for (int i = 0; i < kNumElements; ++i) { if (output[i] != expected[i]) { return absl::InternalError( absl::StrFormat("device %d output[%d] = %g, expected %g", @@ -201,8 +202,8 @@ TEST(CollectivePermuteThunkMultiGpuTest, ExecuteOnStream) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); CollectivePermuteThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -222,8 +223,8 @@ TEST(CollectivePermuteThunkMultiGpuTest, RecordCommandBufferCreate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); CollectivePermuteThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); @@ -243,8 +244,8 @@ TEST(CollectivePermuteThunkMultiGpuTest, RecordCommandBufferUpdate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation alloc_src(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation alloc_dst(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); CollectivePermuteThunk thunk = MakeThunk(alloc_src, alloc_dst); std::vector slots(kNumDevices); diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc new file mode 100644 index 00000000000000..4ae03ebd3441d8 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc @@ -0,0 +1,248 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/runtime/collective_reduce_thunk.h" + +#include +#include +#include +#include +#include + +#include "absl/base/casts.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "xla/backends/gpu/collectives/gpu_clique_key.h" +#include "xla/backends/gpu/collectives/gpu_collectives.h" +#include "xla/backends/gpu/collectives/gpu_communicator.h" +#include "xla/backends/gpu/runtime/all_reduce_thunk.h" +#include "xla/backends/gpu/runtime/collective_thunk.h" +#include "xla/backends/gpu/runtime/collective_thunk.pb.h" +#include "xla/backends/gpu/runtime/thunk.h" +#include "xla/backends/gpu/runtime/thunk.pb.h" +#include "xla/core/collectives/communicator.h" +#include "xla/core/collectives/rank_id.h" +#include "xla/core/collectives/reduction_kind.h" +#include "xla/future.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/service/buffer_assignment.h" +#include "xla/service/collective_ops_utils.h" +#include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/memory_allocation.h" +#include "xla/stream_executor/stream.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tsl/platform/logging.h" +#include "xla/util.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace gpu { + +CollectiveReduceThunk::CollectiveReduceThunk(ThunkInfo thunk_info, + AllReduceConfig config, + std::vector buffers, + bool has_dynamic_root) + : AllReduceReduceScatterThunkBase(Thunk::kCollectiveReduce, thunk_info, + std::move(config), std::move(buffers)), + has_dynamic_root_(has_dynamic_root) {} + +CollectiveReduceThunk::CollectiveReduceThunk( + ThunkInfo thunk_info, const HloCollectiveReduceInstruction* inst, + std::vector buffers, bool p2p_memcpy_enabled, bool has_dynamic_root) + : AllReduceReduceScatterThunkBase(Thunk::kCollectiveReduce, thunk_info, + GetAllReduceConfigInst(inst), + std::move(buffers)), + has_dynamic_root_(has_dynamic_root) {} + +/*static*/ absl::Status CollectiveReduceThunk::CheckImplementable( + const HloCollectiveReduceInstruction* inst, int64_t replica_count, + int64_t partition_count) { + auto status = [&]() -> absl::Status { + // The trailing S32 root operand (dynamic root) is not reduced, so exclude + // it from operand validation. + int64_t num_data_operands = inst->has_dynamic_root() + ? inst->operand_count() - 1 + : inst->operand_count(); + for (int64_t i = 0; i < num_data_operands; ++i) { + ABSL_RETURN_IF_ERROR( + IsValidOperand(inst->operand(i)->shape(), Thunk::kCollectiveReduce)); + } + if (!MatchReductionComputation(inst->called_computations().front()) + .has_value()) { + return absl::UnimplementedError("Unrecognized reduction computation"); + } + return absl::OkStatus(); + }(); + return AddOpDescription(status, inst, replica_count, + partition_count); +} + +/*static*/ CollectiveOpGroupMode CollectiveReduceThunk::GetGroupMode( + const HloCollectiveReduceInstruction* inst) { + return GetAllReduceConfigInst(inst).config.group_mode; +} + +absl::Status CollectiveReduceThunk::InitializeCollective( + const InitializeParams& params, const GpuCliqueKey& clique_key) { + if (!has_dynamic_root_) { + return absl::OkStatus(); + } + se::StreamExecutor* executor = params.executor; + absl::MutexLock lock(mutex_); + std::unique_ptr& metadata = + per_executor_metadata_[executor]; + if (metadata == nullptr) { + metadata = std::make_unique(); + } + if (metadata->reduce_roots == nullptr) { + // The last buffer holds the runtime-selected root ranks (one S32 per + // reduce); all other buffers are the data being reduced. + metadata->num_roots = buffers().size() - 1; + ABSL_ASSIGN_OR_RETURN( + std::unique_ptr alloc, + executor->HostMemoryAllocate(metadata->num_roots * sizeof(int32_t))); + metadata->reduce_roots = std::move(alloc); + } + return absl::OkStatus(); +} + +absl::Status CollectiveReduceThunk::RunCollective( + const ExecuteParams& params, const GpuCliqueKey& clique_key, + se::Stream& stream, Communicator& comm) { + ABSL_ASSIGN_OR_RETURN(std::vector device_buffers, + ConvertToDeviceBuffers(params.buffer_allocations, buffers(), + config_.config.operand_element_type)); + CollectiveReduceMetadata* metadata = nullptr; + { + absl::MutexLock lock(mutex_); + metadata = per_executor_metadata_[stream.parent()].get(); + } + return RunCollectiveReduce(config_.reduction_kind, device_buffers, stream, + comm, metadata, has_dynamic_root_); +} + +absl::StatusOr> +CollectiveReduceThunk::FromProto( + ThunkInfo thunk_info, const CollectiveReduceThunkProto& thunk_proto, + absl::Span buffer_allocations) { + std::vector buffers; + buffers.reserve(thunk_proto.buffers_size()); + for (const CollectiveBufferProto& proto : thunk_proto.buffers()) { + ABSL_ASSIGN_OR_RETURN( + CollectiveThunk::Buffer buffer, + CollectiveThunk::Buffer::FromProto(proto, buffer_allocations)); + buffers.push_back(buffer); + } + + CollectiveConfig config = + CollectiveConfig::FromProto(thunk_proto.collective_config()); + + ABSL_ASSIGN_OR_RETURN(ReductionKind reduction_kind, + FromReductionKindProto(thunk_proto.reduction_kind())); + + return std::make_unique( + std::move(thunk_info), AllReduceConfig{config, reduction_kind}, + std::move(buffers), thunk_proto.has_dynamic_root()); +} + +absl::StatusOr CollectiveReduceThunk::ToProto() const { + ThunkProto proto; + *proto.mutable_thunk_info() = thunk_info().ToProto(); + + CollectiveReduceThunkProto* thunk_proto = + proto.mutable_collective_reduce_thunk(); + + for (const Buffer& buffer : buffers()) { + ABSL_ASSIGN_OR_RETURN(*thunk_proto->add_buffers(), buffer.ToProto()); + } + + *thunk_proto->mutable_collective_config() = config_.config.ToProto(); + thunk_proto->set_reduction_kind(ToReductionKindProto(config_.reduction_kind)); + thunk_proto->set_has_dynamic_root(has_dynamic_root_); + + return proto; +} + +absl::Status RunCollectiveReduce(ReductionKind reduction_kind, + std::vector& buffers, + se::Stream& stream, Communicator& comm, + CollectiveReduceMetadata* metadata, + bool has_dynamic_root) { + int device_ordinal = stream.parent()->device_ordinal(); + XLA_VLOG_DEVICE(3, device_ordinal) << "Performing collective-reduce"; + + if (has_dynamic_root && metadata) { + // Stage the runtime-selected root ranks on the host so we can pick a root + // per reduce below. + DeviceBufferPair& roots_device_buffer = buffers.back(); + CHECK(metadata->reduce_roots != nullptr); + ABSL_RETURN_IF_ERROR(stream.Memcpy(metadata->reduce_roots->address().opaque(), + roots_device_buffer.source_buffer, + roots_device_buffer.source_buffer.size())); + if (absl::Status blocked = stream.BlockHostUntilDone(); !blocked.ok()) { + return absl::InternalError( + absl::StrFormat("Failed to copy dynamic roots on stream %p: %s", + &stream, blocked.message())); + } + } + + // With a dynamic root, the last buffer only carries the per-reduce root ranks + // and must not itself be reduced. + const int64_t num_reduces = + (has_dynamic_root && metadata) ? buffers.size() - 1 : buffers.size(); + auto* gpu_comm = absl::down_cast(&comm); + + if (has_dynamic_root && metadata) { + ABSL_ASSIGN_OR_RETURN(size_t num_ranks, comm.NumRanks()); + int32_t* roots_ptr = + reinterpret_cast(metadata->reduce_roots->address().opaque()); + for (int64_t i = 0; i < num_reduces; ++i) { + if (roots_ptr[i] < 0 || static_cast(roots_ptr[i]) >= num_ranks) { + return absl::InvalidArgumentError(absl::StrFormat( + "collective-reduce dynamic root %d at index %d is out of range " + "[0, %d)", + roots_ptr[i], i, num_ranks)); + } + } + } + + Future<> future = gpu_comm->GroupExecute([&]() -> absl::Status { + RankId root(0); + for (int64_t i = 0; i < num_reduces; ++i) { + DeviceBufferPair& buffer = buffers[i]; + if (has_dynamic_root && metadata) { + int32_t* roots_ptr = reinterpret_cast( + metadata->reduce_roots->address().opaque()); + root = RankId(roots_ptr[i]); + } + ABSL_RETURN_IF_ERROR(gpu_comm->LaunchReduce( + buffer.source_buffer, buffer.destination_buffer, buffer.element_type, + buffer.element_count, reduction_kind, root, + GpuCollectives::On(stream))); + } + return absl::OkStatus(); + }); + ABSL_RETURN_IF_ERROR(future.Await()); + XLA_VLOG_DEVICE(3, device_ordinal) << "Done performing collective-reduce"; + return absl::OkStatus(); +} + +} // namespace gpu +} // namespace xla diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h new file mode 100644 index 00000000000000..c08fcafb559dd7 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h @@ -0,0 +1,113 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_REDUCE_THUNK_H_ +#define XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_REDUCE_THUNK_H_ + +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/container/flat_hash_map.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "xla/backends/gpu/collectives/gpu_clique_key.h" +#include "xla/backends/gpu/runtime/all_reduce_thunk.h" +#include "xla/backends/gpu/runtime/collective_thunk.h" +#include "xla/backends/gpu/runtime/thunk.pb.h" +#include "xla/core/collectives/communicator.h" +#include "xla/core/collectives/reduction_kind.h" +#include "xla/hlo/ir/hlo_instructions.h" +#include "xla/service/buffer_assignment.h" +#include "xla/stream_executor/memory_allocation.h" +#include "xla/stream_executor/stream.h" +#include "xla/stream_executor/stream_executor.h" + +namespace xla { +namespace gpu { + +// Per-`StreamExecutor` state used only when the reduce root is selected at run +// time. `reduce_roots` is a host-pinned staging buffer that the trailing S32 +// root operand is copied into before the reduce is launched. +struct CollectiveReduceMetadata { + int64_t num_roots = 0; + std::unique_ptr reduce_roots = nullptr; +}; + +// Thunk that performs a NCCL-based reduce (`ncclReduce`): the operands are +// reduced with the `to_apply` computation and the result is written only to the +// root rank. The root is rank 0 within the replica group by default, or chosen +// at run time from a trailing S32 operand when `has_dynamic_root` is set. +class CollectiveReduceThunk : public AllReduceReduceScatterThunkBase { + public: + CollectiveReduceThunk(ThunkInfo thunk_info, AllReduceConfig config, + std::vector buffers, + bool has_dynamic_root = false); + CollectiveReduceThunk(ThunkInfo thunk_info, + const HloCollectiveReduceInstruction* inst, + std::vector buffers, + bool p2p_memcpy_enabled = false, + bool has_dynamic_root = false); + + static absl::string_view GetHloOpName() { return "collective-reduce-start"; } + + static absl::Status CheckImplementable( + const HloCollectiveReduceInstruction* inst, int64_t replica_count, + int64_t partition_count); + + static CollectiveOpGroupMode GetGroupMode( + const HloCollectiveReduceInstruction* inst); + + static absl::StatusOr> FromProto( + ThunkInfo thunk_info, const CollectiveReduceThunkProto& thunk_proto, + absl::Span buffer_allocations); + + absl::StatusOr ToProto() const override; + + protected: + bool RequiresRendezvous() const override { return true; } + + absl::Status InitializeCollective(const InitializeParams& params, + const GpuCliqueKey& clique_key) override; + + absl::Status RunCollective(const ExecuteParams& params, + const GpuCliqueKey& clique_key, se::Stream& stream, + Communicator& comm) override; + + private: + const bool has_dynamic_root_; + mutable absl::Mutex mutex_; + absl::flat_hash_map> + per_executor_metadata_ ABSL_GUARDED_BY(mutex_); +}; + +// Runs a reduce over `buffers` on `stream`. When `has_dynamic_root` is set the +// last entry of `buffers` is the S32 root vector (not itself reduced) and +// `metadata` provides the host staging buffer for the root ranks. +absl::Status RunCollectiveReduce(ReductionKind reduction_kind, + std::vector& buffers, + se::Stream& stream, Communicator& comm, + CollectiveReduceMetadata* metadata = nullptr, + bool has_dynamic_root = false); + +} // namespace gpu +} // namespace xla + +#endif // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_REDUCE_THUNK_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc new file mode 100644 index 00000000000000..f903447f705c99 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc @@ -0,0 +1,214 @@ +/* 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. +==============================================================================*/ + +// Multi-GPU integration tests for CollectiveReduceThunk through NCCL +// (ncclReduce). Requires at least kNumDevices GPUs. + +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/strings/str_format.h" +#include "absl/types/span.h" +#include "xla/backends/gpu/runtime/all_reduce_thunk.h" +#include "xla/backends/gpu/runtime/collective_reduce_thunk.h" +#include "xla/backends/gpu/runtime/collective_thunk.h" +#include "xla/backends/gpu/runtime/collective_thunk_multigpu_test_utils.h" +#include "xla/backends/gpu/runtime/thunk.h" +#include "xla/core/collectives/reduction_kind.h" +#include "xla/service/buffer_assignment.h" +#include "xla/service/device_assignment.h" +#include "xla/service/gpu/buffer_allocations.h" +#include "xla/service/shaped_slice.h" +#include "xla/shape_util.h" +#include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/stream.h" +#include "xla/xla_data.pb.h" + +namespace xla::gpu { +namespace { + +static constexpr int kNumDevices = 2; + +static AllReduceConfig MakeCollectiveReduceConfig(bool has_dynamic_root) { + CollectiveConfig config; + config.operand_element_type = {F32}; + if (has_dynamic_root) { + config.operand_element_type.push_back(S32); + } + config.group_mode = COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA; + + ReplicaGroup replica_group; + for (int i = 0; i < kNumDevices; ++i) { + replica_group.add_replica_ids(i); + } + config.replica_groups = {replica_group}; + return AllReduceConfig{config, ReductionKind::SUM}; +} + +// Builds a sum-reduce thunk. When `alloc_root` is provided the reduce root is +// selected at run time from that S32 buffer, otherwise it defaults to rank 0. +static CollectiveReduceThunk MakeThunk( + const BufferAllocation& alloc_src, const BufferAllocation& alloc_dst, + const BufferAllocation* alloc_root = nullptr) { + ShapedSlice src_slice{ + BufferAllocation::Slice(&alloc_src, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + ShapedSlice dst_slice{ + BufferAllocation::Slice(&alloc_dst, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + CollectiveThunk::Buffer buffer{.element_count = kNumElements, + .source_buffer = src_slice, + .destination_buffer = dst_slice, + .source_memory_space = 0, + .destination_memory_space = 0}; + std::vector buffers = {buffer}; + if (alloc_root != nullptr) { + ShapedSlice root_slice{ + BufferAllocation::Slice(alloc_root, 0, sizeof(int32_t)), + ShapeUtil::MakeShape(S32, {1})}; + buffers.push_back(CollectiveThunk::Buffer{.element_count = 1, + .source_buffer = root_slice, + .destination_buffer = root_slice, + .source_memory_space = 0, + .destination_memory_space = 0}); + } + const bool has_dynamic_root = alloc_root != nullptr; + return CollectiveReduceThunk(Thunk::ThunkInfo(), + MakeCollectiveReduceConfig(has_dynamic_root), + buffers, has_dynamic_root); +} + +using DeviceTestSlot = CollectiveThunkMultiGpuTestState; + +static std::vector SourceValues(int device_ordinal) { + std::vector values(kNumElements); + for (int i = 0; i < kNumElements; ++i) { + values[i] = static_cast(device_ordinal * 10 + i); + } + return values; +} + +// The reduced (summed) value expected on the root rank. +static std::vector ExpectedSum() { + std::vector expected(kNumElements, 0.0f); + for (int d = 0; d < kNumDevices; ++d) { + std::vector values = SourceValues(d); + for (int i = 0; i < kNumElements; ++i) { + expected[i] += values[i]; + } + } + return expected; +} + +static absl::Status VerifyRootOutput(se::Stream& stream, + se::DeviceAddressBase dst, + int device_ordinal) { + ABSL_ASSIGN_OR_RETURN(std::vector output, + ReadDeviceBuffer(stream, dst, kNumElements)); + std::vector expected = ExpectedSum(); + for (int i = 0; i < kNumElements; ++i) { + if (output[i] != expected[i]) { + return absl::InternalError( + absl::StrFormat("root device %d output[%d] = %g, expected %g", + device_ordinal, i, output[i], expected[i])); + } + } + return absl::OkStatus(); +} + +// Runs the reduce on each device and verifies the sum landed on `root_rank`. +static absl::Status RunReduce(std::vector& slots, + CollectiveReduceThunk& thunk, + const DeviceAssignment& device_assignment, + int64_t root_buffer_size, int root_rank) { + return RunOnDevices( + kNumDevices, "collective_reduce_execute", [&](int d) -> absl::Status { + std::vector buffer_sizes = {kFloatByteLength, + kFloatByteLength}; + if (root_buffer_size > 0) { + buffer_sizes.push_back(root_buffer_size); + } + ABSL_RETURN_IF_ERROR(SetupCollectiveThunkDevice( + d, kNumDevices, buffer_sizes, thunk, device_assignment, slots[d])); + + DeviceTestSlot& slot = slots[d]; + ABSL_RETURN_IF_ERROR(FillDeviceBuffer(*slot.stream, slot.create_buffers[0], + SourceValues(d))); + ABSL_RETURN_IF_ERROR( + FillDeviceBuffer(*slot.stream, slot.create_buffers[1], + std::vector(kNumElements, -1.0f))); + int32_t root_value = static_cast(root_rank); + if (root_buffer_size > 0) { + se::DeviceAddressBase root_buf = slot.create_buffers[2]; + ABSL_RETURN_IF_ERROR( + slot.stream->Memcpy(&root_buf, &root_value, sizeof(int32_t))); + // Block so `root_value` outlives the asynchronous host-to-device + // copy. + ABSL_RETURN_IF_ERROR(slot.stream->BlockHostUntilDone()); + } + + BufferAllocations allocations = + MakeBufferAllocations(slot, slot.create_buffers); + Thunk::ExecuteParams execute_params = + MakeExecuteParams(slot, allocations); + ABSL_RETURN_IF_ERROR(ExecuteOnStreamAndBlock(thunk, execute_params)); + + // The reduce writes the result only to the root rank; other ranks' + // destination buffers are left undefined. + if (d == root_rank) { + return VerifyRootOutput(*slot.stream, slot.create_buffers[1], d); + } + return absl::OkStatus(); + }); +} + +TEST(CollectiveReduceThunkMultiGpuTest, ExecuteOnStream) { + if (!HasEnoughGpus(kNumDevices)) { + GTEST_SKIP() << "Test requires at least " << kNumDevices << " GPUs"; + } + + DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); + CollectiveReduceThunk thunk = MakeThunk(alloc_src, alloc_dst); + std::vector slots(kNumDevices); + + ASSERT_OK(RunReduce(slots, thunk, device_assignment, + /*root_buffer_size=*/0, /*root_rank=*/0)); +} + +TEST(CollectiveReduceThunkMultiGpuTest, ExecuteOnStreamWithDynamicRoot) { + if (!HasEnoughGpus(kNumDevices)) { + GTEST_SKIP() << "Test requires at least " << kNumDevices << " GPUs"; + } + + DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); + BufferAllocation alloc_src(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_dst(/*index=*/1, kFloatByteLength, /*color=*/0); + BufferAllocation alloc_root(/*index=*/2, sizeof(int32_t), /*color=*/0); + CollectiveReduceThunk thunk = MakeThunk(alloc_src, alloc_dst, &alloc_root); + std::vector slots(kNumDevices); + + // Select rank 1 as the root at run time; the sum must land on device 1. + ASSERT_OK(RunReduce(slots, thunk, device_assignment, + /*root_buffer_size=*/alloc_root.size(), /*root_rank=*/1)); +} + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc index 0ad331a60d8761..e0762966398a04 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc @@ -326,11 +326,11 @@ absl::StatusOr> GetCollectiveBuffers( } else { // For other operations simply zip operands with results. // - // A collective-broadcast with a dynamic root carries an extra trailing - // operand: a 1-D S32 vector holding the runtime-selected root rank for each - // data operand. That operand has no corresponding output, so it is not - // zipped with a result; instead it is mapped to its own allocation and - // consumed separately by the thunk. + // A collective-broadcast or collective-reduce with a dynamic root carries + // an extra trailing operand: a 1-D S32 vector holding the runtime-selected + // root rank for each data operand. That operand has no corresponding + // output, so it is not zipped with a result; instead it is mapped to its + // own allocation and consumed separately by the thunk. int64_t num_data_operands = has_dynamic_root ? operand_count - 1 : operand_count; for (int64_t i = 0; i < num_data_operands; i++) { diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_thunk_multigpu_test_utils.h b/third_party/xla/xla/backends/gpu/runtime/collective_thunk_multigpu_test_utils.h index afa27ccbbb78e0..d3684b39c34a79 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_thunk_multigpu_test_utils.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_thunk_multigpu_test_utils.h @@ -48,6 +48,9 @@ limitations under the License. namespace xla::gpu { +inline constexpr int64_t kNumElements = 4; +inline constexpr int64_t kFloatByteLength = sizeof(float) * kNumElements; + struct CollectiveThunkMultiGpuTestState { int device_ordinal = 0; se::StreamExecutor* executor = nullptr; diff --git a/third_party/xla/xla/backends/gpu/runtime/command.h b/third_party/xla/xla/backends/gpu/runtime/command.h index 5bf712cdfc0b61..c47820ae90e72f 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command.h +++ b/third_party/xla/xla/backends/gpu/runtime/command.h @@ -25,7 +25,6 @@ limitations under the License. #include #include "absl/functional/function_ref.h" -#include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -241,36 +240,6 @@ Command::WalkResult Command::Walk(F&& callback) const { }); } -//===----------------------------------------------------------------------===// -// Asynchronous commands -//===----------------------------------------------------------------------===// - -// A base class for a command that starts an asynchronous execution. -class AsyncStartCommand : public Command { - public: - using Command::Command; - - // At run time async command might behave like a synchronous one, i.e. - // some collective operations if they can't be overlapped with compute - // operations executed like they have synchronous execution semantics. - virtual bool IsAsync() const = 0; -}; - -// A command that completes an `async_start` command. -class AsyncDoneCommand : public Command { - public: - explicit AsyncDoneCommand(const AsyncStartCommand* async_start) - : Command(Thunk::Kind::kAsyncDone), async_start_(async_start) { - DCHECK(async_start_) << "AsyncStart command must be not null"; - } - - const AsyncStartCommand* async_start() const { return async_start_; } - bool IsAsync() const { return async_start_->IsAsync(); } - - private: - const AsyncStartCommand* async_start_; -}; - //===----------------------------------------------------------------------===// // CommandSequence //===----------------------------------------------------------------------===// diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter.cc index 2568cb0633ce24..d0dbc4a5a31e6a 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter.cc @@ -57,6 +57,26 @@ namespace { // A context for tracking thunks to commands conversion details. struct ConversionContext { std::vector extra_resources; + + void Append(CommandSequence& commands, Command* command, + const ConvertToCommandsOptions& options) { + if (options.synchronization_mode == + CommandExecutor::SynchronizationMode::kLHS) { + // Preserve the existing serial order of the flattened command sequence. + // The executor infers this order from token dependencies. + // + // TODO(shawnwang18): This serial chain is temporary. It serializes + // flattened AsyncStartThunk bodies with the main stream. Replace it with + // fork/join dependencies for AsyncStartThunk / AsyncDoneThunk in + // https://github.com/openxla/xla/pull/48900. + Command::ResourceUses dependencies; + if (!commands.empty()) { + dependencies.push_back(ResourceUse::Read(commands.back()->token())); + } + extra_resources.push_back(std::move(dependencies)); + } + commands.Append(command); + } }; // A thunk with its concurrent region id in the case where inherited @@ -162,7 +182,7 @@ static absl::Status AppendCommands(ConversionContext& ctx, auto& conditional_thunk = static_cast(thunk); ABSL_RETURN_IF_ERROR( SetOrUpdateCommandBufferBranchExecutors(conditional_thunk, options)); - cmd_sequence.Append(&conditional_thunk); + ctx.Append(cmd_sequence, &conditional_thunk, options); return absl::OkStatus(); } case Thunk::Kind::kAsyncDone: @@ -171,7 +191,7 @@ static absl::Status AppendCommands(ConversionContext& ctx, case Thunk::Kind::kWhile: { auto& while_thunk = static_cast(thunk); ABSL_RETURN_IF_ERROR(SetOrUpdateCommandBufferExecutors(while_thunk, options)); - cmd_sequence.Append(&while_thunk); + ctx.Append(cmd_sequence, &while_thunk, options); return absl::OkStatus(); } case Thunk::Kind::kDynamicSliceFusion: { @@ -183,7 +203,7 @@ static absl::Status AppendCommands(ConversionContext& ctx, ABSL_RETURN_IF_ERROR( dynamic_slice_fusion_thunk.SetOrUpdateCommandBufferExecutor( std::move(cmds))); - cmd_sequence.Append(&dynamic_slice_fusion_thunk); + ctx.Append(cmd_sequence, &dynamic_slice_fusion_thunk, options); return absl::OkStatus(); } // Sequential thunk does not have any special semantics and we simply inline @@ -213,7 +233,7 @@ static absl::Status AppendCommands(ConversionContext& ctx, if (auto* command = dynamic_cast(&thunk)) { // Command/thunk hybrids are owned by the input ThunkSequence and outlive // the returned CommandSequence, so command sequences borrow them directly. - cmd_sequence.Append(command); + ctx.Append(cmd_sequence, command, options); return absl::OkStatus(); } diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc index 12aae7ff1e77a0..512ceb40b57405 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc @@ -27,9 +27,11 @@ limitations under the License. #include "absl/status/status.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "xla/backends/gpu/runtime/async_thunk.h" #include "xla/backends/gpu/runtime/command.h" #include "xla/backends/gpu/runtime/command_executor.h" #include "xla/backends/gpu/runtime/conditional_thunk.h" +#include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/runtime/kernel_thunk.h" #include "xla/backends/gpu/runtime/sequential_thunk.h" #include "xla/backends/gpu/runtime/thunk.h" @@ -38,13 +40,17 @@ limitations under the License. #include "xla/codegen/emitters/kernel_arguments.h" #include "xla/runtime/execution_graph.h" #include "xla/service/buffer_assignment.h" +#include "xla/service/gpu/launch_dimensions.h" #include "xla/service/shaped_slice.h" #include "xla/shape_util.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/stream_executor/gpu/tma_metadata.h" +#include "xla/stream_executor/launch_dim.h" #include "xla/tsl/platform/test.h" #include "xla/xla_data.pb.h" +#include "tsl/platform/status_matchers.h" namespace xla::gpu { +namespace { using ::testing::AllOf; using ::testing::ElementsAre; @@ -855,4 +861,101 @@ TEST_F(CommandBufferCmdEmitterTest, ConcurrentRegionsForwardDependencies) { EXPECT_THAT(nodes[name_to_ids["d"]].out_edges, IsEmpty()); } +class AsyncCommandBufferCmdEmitterTest : public CommandBufferCmdEmitterTest { + protected: + void Kernel(ThunkSequence& thunks, absl::string_view name, int64_t slot) { + thunks.Emplace( + NextThunkInfo(name), + BufferAllocation::Slice(&allocation_, slot * 1024, 1024)); + } + + AsyncStartThunk* Start(ThunkSequence& thunks, ExecutionStreamId stream, + ThunkSequence body) { + auto start = std::make_unique(NextThunkInfo("start"), + stream, std::move(body)); + AsyncStartThunk* result = start.get(); + thunks.push_back(std::move(start)); + return result; + } + + void Done(ThunkSequence& thunks, const AsyncStartThunk* start, + absl::string_view name) { + thunks.Emplace(NextThunkInfo(name), + start->async_execution()); + } + + // Check reachability rather than direct edges: graph construction can remove + // redundant edges without changing the execution order. + bool HappensBefore(CommandExecutor& commands, const std::string& before, + const std::string& after) { + auto ids = NamesToNodeIds(commands); + if (ids.find(before) == ids.end() || ids.find(after) == ids.end()) { + return false; + } + auto graph = commands.execution_graph(); + std::vector reachable(commands.size(), false); + reachable[ids.at(before)] = true; + for (int64_t i = ids.at(before); i < static_cast(commands.size()); + ++i) { + if (!reachable[i]) { + continue; + } + for (const auto& edge : graph->nodes_defs()[i].out_edges) { + reachable[edge.id] = true; + } + } + return reachable[ids.at(after)]; + } + + const ConvertToCommandsOptions options_ = { + CommandExecutor::SynchronizationMode::kLHS}; + + private: + BufferAllocation allocation_{0, 16 * 1024, 0}; +}; + +TEST_F(AsyncCommandBufferCmdEmitterTest, KeepsOrdinaryCommandsOrdered) { + ThunkSequence thunks; + Kernel(thunks, "first", 0); + ThunkSequence nested; + Kernel(nested, "second", 1); + thunks.Emplace(NextThunkInfo("sequential"), + std::move(nested)); + Kernel(thunks, "third", 2); + + ASSERT_OK_AND_ASSIGN(CommandExecutor commands, + ConvertToCommands(thunks, options_)); + ASSERT_EQ(commands.size(), 3); + EXPECT_TRUE(commands.execution_graph()->is_sequential()); + EXPECT_TRUE(HappensBefore(commands, "first", "second")); + EXPECT_TRUE(HappensBefore(commands, "second", "third")); +} + +TEST_F(AsyncCommandBufferCmdEmitterTest, PreservesFlattenedAsyncOrder) { + ThunkSequence thunks; + Kernel(thunks, "before", 0); + ThunkSequence body; + Kernel(body, "async_a", 1); + ThunkSequence nested; + Kernel(nested, "async_b", 2); + body.Emplace(NextThunkInfo("sequential"), std::move(nested)); + AsyncStartThunk* start = + Start(thunks, CommunicationStreamId(0), std::move(body)); + Kernel(thunks, "main_a", 3); + Kernel(thunks, "main_b", 4); + Done(thunks, start, "done"); + Kernel(thunks, "after", 5); + + ASSERT_OK_AND_ASSIGN(CommandExecutor commands, + ConvertToCommands(thunks, options_)); + EXPECT_EQ(commands.size(), 6); + EXPECT_TRUE(commands.execution_graph()->is_sequential()); + EXPECT_TRUE(HappensBefore(commands, "before", "async_a")); + EXPECT_TRUE(HappensBefore(commands, "async_a", "async_b")); + EXPECT_TRUE(HappensBefore(commands, "async_b", "main_a")); + EXPECT_TRUE(HappensBefore(commands, "main_a", "main_b")); + EXPECT_TRUE(HappensBefore(commands, "main_b", "after")); +} + +} // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc index 295772501d1968..ba4312fcbaafb5 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc @@ -86,6 +86,8 @@ std::optional GetCollectiveOpType( return DebugOptions::COLLECTIVEBROADCAST; case Thunk::kCollectivePermute: return DebugOptions::COLLECTIVEPERMUTE; + case Thunk::kCollectiveReduce: + return DebugOptions::ALLREDUCE; case Thunk::kRaggedAllToAll: return DebugOptions::RAGGEDALLTOALL; case Thunk::kReduceScatter: @@ -212,6 +214,7 @@ std::optional GetCommandBufferCmdType( case Thunk::kAllToAll: case Thunk::kCollectiveBroadcast: case Thunk::kCollectivePermute: + case Thunk::kCollectiveReduce: case Thunk::kGroup: case Thunk::kRaggedAllToAll: case Thunk::kReduceScatter: diff --git a/third_party/xla/xla/backends/gpu/runtime/command_executor.cc b/third_party/xla/xla/backends/gpu/runtime/command_executor.cc index fa705edf7f5bfe..966e7b84c04380 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_executor.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_executor.cc @@ -137,98 +137,6 @@ std::vector CreateCommandOperationsWithConcurrentMode( return operations; } -// Helper: Check if a command is an Async Start -bool IsAsyncStart(const Command* cmd) { - const auto* async = dynamic_cast(cmd); - return async && async->IsAsync(); -} - -// Helper: Check if a command is an Async Done -bool IsAsyncDone(const Command* cmd) { - const auto* async = dynamic_cast(cmd); - return async && async->IsAsync(); -} - -// Helper: Find the corresponding Start command for a given Done command -int64_t FindMatchingStartId(const CommandSequence& commands, int64_t done_idx) { - const auto* done_cmd = - dynamic_cast(commands[done_idx]); - CHECK(done_cmd); - - for (int64_t j = done_idx - 1; j >= 0; --j) { - if (!IsAsyncStart(commands[j])) { - continue; - } - - const auto* start_cmd = dynamic_cast(commands[j]); - CHECK(start_cmd); - - if (start_cmd->IsAsync() && done_cmd->async_start() == start_cmd) { - return j; - } - } - return -1; -} - -// Helper: Add dependency on the nearest previous command that is NOT an Async -// Start, by pushing into `extras`. -void AddDependencyOnPrevNonStart(const CommandSequence& commands, - int64_t current_idx, - Command::ResourceUses& extras) { - for (int64_t j = current_idx - 1; j >= 0; --j) { - if (IsAsyncStart(commands[j])) { - // Skip other starts - continue; - } - - extras.push_back(ResourceUse::Read(commands[j]->token())); - break; // Found the dependency, stop scanning - } -} - -std::vector CreateCommandOperationsWithLHSMode( - const CommandSequence& commands, - absl::Span extra_resources) { - VLOG(3) << "CreateCommandOperations with synchronization mode: LHS"; - - // 1. Dependency Analysis Phase: pre-compute LHS resource uses per command. - std::vector lhs_extras(commands.size()); - for (int64_t i = 0; i < static_cast(commands.size()); ++i) { - if (IsAsyncDone(commands[i])) { - // CASE A: Async Done — depends on its matching Start command - int64_t start_id = FindMatchingStartId(commands, i); - CHECK_NE(start_id, -1); - lhs_extras[i].push_back(ResourceUse::Read(commands[start_id]->token())); - - // Also depends on immediate predecessor (if it's not the start itself) - CHECK_GT(i, 0); - if ((i - 1) != start_id) { - lhs_extras[i].push_back(ResourceUse::Read(commands[i - 1]->token())); - } - } else { - // CASE B: Standard Command OR Async Start - // Both share the same logic: depend on the previous non-async-start - AddDependencyOnPrevNonStart(commands, i, lhs_extras[i]); - } - } - - // 2. Construction Phase: build operations with merged extra resources. - std::vector operations; - operations.reserve(commands.size()); - for (size_t i = 0; i < commands.size(); ++i) { - Command::ResourceUses merged; - if (!extra_resources.empty()) { - merged.insert(merged.end(), extra_resources[i].begin(), - extra_resources[i].end()); - } - merged.insert(merged.end(), lhs_extras[i].begin(), lhs_extras[i].end()); - operations.emplace_back(commands[i], merged); - } - - VlogOperations(operations); - return operations; -} - } // namespace static absl::StatusOr> CreateCommandOperations( @@ -237,18 +145,17 @@ static absl::StatusOr> CreateCommandOperations( absl::Span extra_resources) { using Mode = CommandExecutor::SynchronizationMode; switch (synchronization_mode) { - // Building an execution graph works the same for kConcurrent and - // kConcurrentRegions. + // Building an execution graph works the same for kConcurrent, + // kConcurrentRegions and kLHS: dependencies are inferred from buffer + // conflicts and from the resources in `extra_resources`. In kLHS mode the + // thunk schedule order is expressed by the emitter as token resources. case Mode::kConcurrent: - case Mode::kConcurrentRegions: { + case Mode::kConcurrentRegions: + case Mode::kLHS: { return CreateCommandOperationsWithConcurrentMode(commands, extra_resources); } - case Mode::kLHS: { - return CreateCommandOperationsWithLHSMode(commands, extra_resources); - } - case Mode::kSerialize: default: { return Unimplemented( diff --git a/third_party/xla/xla/backends/gpu/runtime/command_executor.h b/third_party/xla/xla/backends/gpu/runtime/command_executor.h index c6cfb64b56ed6e..fe2b1edd97d9b7 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_executor.h +++ b/third_party/xla/xla/backends/gpu/runtime/command_executor.h @@ -65,8 +65,9 @@ class CommandExecutor { // that have buffer of resource conflicts, and building a DAG of commands. kConcurrent, - // Uses the same latency hidden scheduling results used in the thunk - // scheduling. + // Preserves the thunk schedule produced by the latency hiding scheduler. + // The emitter expresses command order as token resources in + // `extra_resources`; the execution graph is built like kConcurrent. kLHS, // Hybrid of kSerialize and kConcurrent. Builds a DAG for concurrent diff --git a/third_party/xla/xla/backends/gpu/runtime/command_executor_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_executor_test.cc index 78096393e523de..25ef8636ab0b5a 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_executor_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_executor_test.cc @@ -403,9 +403,9 @@ TEST_F(CommandExecutorRendererTest, RenderUsesStoredExtraResourcesLHSMode) { ASSERT_EQ(fake_renderer_->captured().size(), 2); // cmd0: Write(cmd0->token()) + Write(shared_resource) = 2. EXPECT_EQ(fake_renderer_->captured()[0].resource_use_count, 2); - // cmd1: Write(cmd1->token()) + Read(shared_resource) [extra] - // + Read(cmd0->token()) [lhs_extras: sequential dependency] = 3. - EXPECT_EQ(fake_renderer_->captured()[1].resource_use_count, 3); + // cmd1: Write(cmd1->token()) + Read(shared_resource) = 2. The executor adds + // no LHS-specific dependencies; schedule order comes from `extra_resources`. + EXPECT_EQ(fake_renderer_->captured()[1].resource_use_count, 2); } // Calling RenderExecutionGraph twice produces consistent results. diff --git a/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_multigpu_test.cc index b990e2c10b6204..b0f78767e17f7e 100644 --- a/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_multigpu_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_multigpu_test.cc @@ -47,8 +47,6 @@ namespace xla::gpu { namespace { static constexpr int kNumDevices = 2; -static constexpr int64_t kLength = 4; -static constexpr int64_t kByteLength = sizeof(float) * kLength; static P2PConfig MakeSendRecvConfig() { ReplicaGroup replica_group; @@ -66,9 +64,9 @@ static P2PConfig MakeSendRecvConfig() { } static CollectiveThunk::Buffer MakeBuffer(const BufferAllocation& allocation) { - ShapedSlice slice{BufferAllocation::Slice(&allocation, 0, kByteLength), - ShapeUtil::MakeShape(F32, {kLength})}; - return CollectiveThunk::Buffer{.element_count = kLength, + ShapedSlice slice{BufferAllocation::Slice(&allocation, 0, kFloatByteLength), + ShapeUtil::MakeShape(F32, {kNumElements})}; + return CollectiveThunk::Buffer{.element_count = kNumElements, .source_buffer = slice, .destination_buffer = slice, .source_memory_space = 0, @@ -91,12 +89,12 @@ struct DeviceTestSlot : CollectiveThunkMultiGpuTestState { }; static std::vector DeviceBufferSizes() { - return {kByteLength, kByteLength}; + return {kFloatByteLength, kFloatByteLength}; } static std::vector SourceValues(int device_ordinal, int phase) { - std::vector data(kLength); - for (int i = 0; i < kLength; ++i) { + std::vector data(kNumElements); + for (int i = 0; i < kNumElements; ++i) { data[i] = static_cast(phase * 100 + device_ordinal * 10 + i); } return data; @@ -106,13 +104,14 @@ static std::vector ExpectedRecvValues(int device_ordinal, int phase) { if (device_ordinal == 1) { return SourceValues(/*device_ordinal=*/0, phase); } - return std::vector(kLength, 0.0f); + return std::vector(kNumElements, 0.0f); } static absl::Status FillDeviceBufferWithValue(se::Stream& stream, se::DeviceAddressBase buffer, float value) { - return FillDeviceBuffer(stream, buffer, std::vector(kLength, value)); + return FillDeviceBuffer(stream, buffer, + std::vector(kNumElements, value)); } static absl::Status PreparePhaseInputs( @@ -127,9 +126,9 @@ static absl::Status VerifyRecvOutput(DeviceTestSlot& slot, int device_ordinal, int phase, se::DeviceAddressBase recv_dst) { ABSL_ASSIGN_OR_RETURN(std::vector output, - ReadDeviceBuffer(*slot.stream, recv_dst, kLength)); + ReadDeviceBuffer(*slot.stream, recv_dst, kNumElements)); std::vector expected = ExpectedRecvValues(device_ordinal, phase); - for (int i = 0; i < kLength; ++i) { + for (int i = 0; i < kNumElements; ++i) { if (output[i] != expected[i]) { return absl::InternalError( absl::StrFormat("device %d output[%d] = %g, expected %g", @@ -253,8 +252,8 @@ TEST(SendRecvThunkMultiGpuTest, ExecuteOnStream) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation send_alloc(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation recv_alloc(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation send_alloc(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation recv_alloc(/*index=*/1, kFloatByteLength, /*color=*/0); SendThunk send_thunk = MakeSendThunk(send_alloc); RecvThunk recv_thunk = MakeRecvThunk(recv_alloc); std::vector slots(kNumDevices); @@ -276,8 +275,8 @@ TEST(SendRecvThunkMultiGpuTest, RecordCommandBufferCreate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation send_alloc(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation recv_alloc(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation send_alloc(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation recv_alloc(/*index=*/1, kFloatByteLength, /*color=*/0); SendThunk send_thunk = MakeSendThunk(send_alloc); RecvThunk recv_thunk = MakeRecvThunk(recv_alloc); std::vector slots(kNumDevices); @@ -299,8 +298,8 @@ TEST(SendRecvThunkMultiGpuTest, RecordCommandBufferUpdate) { } DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices); - BufferAllocation send_alloc(/*index=*/0, kByteLength, /*color=*/0); - BufferAllocation recv_alloc(/*index=*/1, kByteLength, /*color=*/0); + BufferAllocation send_alloc(/*index=*/0, kFloatByteLength, /*color=*/0); + BufferAllocation recv_alloc(/*index=*/1, kFloatByteLength, /*color=*/0); SendThunk send_thunk = MakeSendThunk(send_alloc); RecvThunk recv_thunk = MakeRecvThunk(recv_alloc); std::vector slots(kNumDevices); diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk.cc b/third_party/xla/xla/backends/gpu/runtime/thunk.cc index 09c67e592375d8..4d07bf0fad704e 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/thunk.cc @@ -158,6 +158,8 @@ ThunkKindProto Thunk::KindToProto(Kind kind) { return THUNK_KIND_BUFFERS_DEBUG_FLOAT_CHECK; case kCollectiveBroadcast: return THUNK_KIND_COLLECTIVE_BROADCAST; + case kCollectiveReduce: + return THUNK_KIND_COLLECTIVE_REDUCE; case kCollectiveKernel: return THUNK_KIND_COLLECTIVE_KERNEL; case kCollectiveMetadata: @@ -259,6 +261,8 @@ absl::StatusOr Thunk::KindFromProto(ThunkKindProto kind) { return kBuffersDebugFloatCheck; case THUNK_KIND_COLLECTIVE_BROADCAST: return kCollectiveBroadcast; + case THUNK_KIND_COLLECTIVE_REDUCE: + return kCollectiveReduce; case THUNK_KIND_COLLECTIVE_KERNEL: return kCollectiveKernel; case THUNK_KIND_COLLECTIVE_METADATA: @@ -359,6 +363,7 @@ absl::StatusOr Thunk::KindFromProto(ThunkKindProto kind) { CASE(kCollectiveKernel); CASE(kCollectiveMetadata); CASE(kCollectivePermute); + CASE(kCollectiveReduce); CASE(kCommand); CASE(kCommandBuffer); CASE(kConditional); @@ -405,7 +410,8 @@ std::ostream& operator<<(std::ostream& os, Thunk::Kind kind) { } bool IsReductionCollective(Thunk::Kind kind) { - return kind == Thunk::kAllReduce || kind == Thunk::kReduceScatter; + return kind == Thunk::kAllReduce || kind == Thunk::kReduceScatter || + kind == Thunk::kCollectiveReduce; } absl::StatusOr Thunk::ThunkInfo::FromProto( @@ -436,6 +442,7 @@ bool Thunk::IsCollective() const { case kAllToAll: case kCollectiveBroadcast: case kCollectivePermute: + case kCollectiveReduce: case kGroup: case kRaggedAllToAll: case kRecv: diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk.h b/third_party/xla/xla/backends/gpu/runtime/thunk.h index 0dce6cff474c5b..971ea903e90e3b 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/thunk.h @@ -108,6 +108,7 @@ class Thunk { kCollectiveKernel, kCollectiveMetadata, kCollectivePermute, + kCollectiveReduce, kCommand, kCommandBuffer, kConditional, diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk.proto b/third_party/xla/xla/backends/gpu/runtime/thunk.proto index 00c062783c010d..231f42399bb217 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk.proto +++ b/third_party/xla/xla/backends/gpu/runtime/thunk.proto @@ -387,6 +387,15 @@ message ReduceScatterThunkProto { xla.ReductionKindProto reduction_kind = 4; } +message CollectiveReduceThunkProto { + CollectiveConfigProto collective_config = 1; + repeated CollectiveBufferProto buffers = 2; + xla.ReductionKindProto reduction_kind = 3; + // When true, the last buffer holds the runtime-selected root rank per reduce + // and is not itself reduced. + bool has_dynamic_root = 4; +} + message AllToAllThunkProto { reserved 1; // Was async_events_unique_id. CollectiveConfigProto collective_config = 2; @@ -521,6 +530,7 @@ message ThunkProto { CollectiveKernelThunkProto collective_kernel_thunk = 59; DynamicSliceFusionThunkProto dynamic_slice_fusion_thunk = 60; RngSeedThunkProto rng_seed_thunk = 61; + CollectiveReduceThunkProto collective_reduce_thunk = 62; } reserved 10, 18, 37, 48, 49, 50, 51, 52, 53, 54, 55, 56; diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_kind.proto b/third_party/xla/xla/backends/gpu/runtime/thunk_kind.proto index 4e0b06f1993aae..a0dd21552b1171 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk_kind.proto +++ b/third_party/xla/xla/backends/gpu/runtime/thunk_kind.proto @@ -32,6 +32,7 @@ enum ThunkKindProto { THUNK_KIND_COLLECTIVE_KERNEL = 15; THUNK_KIND_COLLECTIVE_METADATA = 16; THUNK_KIND_COLLECTIVE_PERMUTE = 17; + THUNK_KIND_COLLECTIVE_REDUCE = 79; THUNK_KIND_COMMAND_BUFFER = 20; THUNK_KIND_CONDITIONAL = 21; THUNK_KIND_CONVOLUTION = 22; diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc index 21d4d9e6bd92c5..0c395bb349ee31 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc +++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc @@ -39,6 +39,7 @@ limitations under the License. #include "xla/backends/gpu/runtime/collective_group_thunk.h" #include "xla/backends/gpu/runtime/collective_kernel_thunk.h" #include "xla/backends/gpu/runtime/collective_permute_thunk.h" +#include "xla/backends/gpu/runtime/collective_reduce_thunk.h" #include "xla/backends/gpu/runtime/conditional_thunk.h" #include "xla/backends/gpu/runtime/convolution_reorder_thunk.h" #include "xla/backends/gpu/runtime/convolution_thunk.h" @@ -301,6 +302,10 @@ absl::StatusOr> DeserializeThunkProtoImpl( return CollectiveBroadcastThunk::FromProto( std::move(thunk_info), thunk_proto.collective_broadcast_thunk(), buffer_allocations); + case ThunkProto::kCollectiveReduceThunk: + return CollectiveReduceThunk::FromProto( + std::move(thunk_info), thunk_proto.collective_reduce_thunk(), + buffer_allocations); case ThunkProto::kCollectiveGroupThunk: return CollectiveGroupThunk::FromProto( std::move(thunk_info), thunk_proto.collective_group_thunk(), diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 9775a8ccd7691f..1f14685f3ec03d 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -763,7 +763,6 @@ lit_test_suite_for_gpus( "dot_bf16.hlo", "element_wise_row_vectorization.hlo", "scatter_bf16.hlo", - "single_instruction.hlo", "reduce_unnested.hlo", "reduction_vectorization_sm_all.hlo", "pdl.hlo", @@ -777,13 +776,11 @@ lit_test_suite_for_gpus( ], "mi350": [ "dot_bf16.hlo", - "single_instruction.hlo", "pdl.hlo", ], "mi450": [ # TODO: re-enable once mi450 memory_bandwidth is finalized. "memory_bandwidth_fusion.hlo", - "single_instruction.hlo", "vectorization.hlo", "pdl.hlo", ], diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc index 15fec5e22fed50..f59c95a6dc0b29 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc @@ -192,6 +192,14 @@ class AsyncCollectiveOps /*memory_size=*/8 * kGB, /*collectives_memory_size=*/std::get<1>(GetParam()) ? 8 * kGB : 0) { } + + protected: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = + CollectiveOpsWithFlagsBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_emit_collective_reduce(true); + return debug_options; + } }; class MemcpyCollectiveOps : public CollectiveOpsWithFlagsBase, @@ -550,6 +558,113 @@ TEST_P(AsyncCollectiveOps, AsyncCollectiveBroadcastDynamicRoot) { } } +TEST_P(AsyncCollectiveOps, AsyncCollectiveReduce) { + // collective-reduce reduces across the replica group and writes the result + // only to the root rank (rank 0 of the group by default), unlike all-reduce. + const absl::string_view kModuleStr = R"( + HloModule test + apply_op { + x = u32[] parameter(0) + y = u32[] parameter(1) + ROOT apply_op = u32[] add(x, y) + } + ENTRY test_computation { + id = u32[] replica-id() + one = u32[] constant(1) + val = u32[] add(id, one) + ROOT cr = u32[] collective-reduce(val), replica_groups={{0, 1}}, + to_apply=apply_op + } + )"; + const int64_t kNumReplicas = 2; + ASSERT_GE(device_count(), kNumReplicas) + << "Test requires at least " << kNumReplicas << " devices (" + << device_count() << " available)"; + + TF_ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + TF_ASSERT_OK_AND_ASSIGN(ExecutionResult execution_result, + ExecuteReplicated(std::move(module))); + + const HloModule* hlo_module = execution_result.optimized_module; + if (enable_async_) { + const HloInstruction* cr_start = + FindCollectiveStart(hlo_module, HloOpcode::kCollectiveReduce); + const HloInstruction* cr_done = + FindCollectiveDone(hlo_module, HloOpcode::kCollectiveReduce); + EXPECT_THAT(cr_start, NotNull()); + EXPECT_THAT(cr_done, NotNull()); + } else { + EXPECT_THAT(FindInstruction(hlo_module, HloOpcode::kCollectiveReduce), + NotNull()); + } + + const std::vector& results = execution_result.results; + ASSERT_EQ(results.size(), kNumReplicas); + // Sum of (replica_id + 1) over replicas {0, 1} = 1 + 2 = 3, written to the + // root rank (replica 0). Non-root outputs are left undefined by the reduce. + LiteralTestUtil::ExpectR0Equal(3, results[0]); +} + +TEST_P(AsyncCollectiveOps, AsyncCollectiveReduceDynamicRoot) { + // The reduce root is supplied at run time by the trailing S32 operand. With + // `replica_groups={{0, 1}}` a static reduce would write to replica 0, so + // selecting root rank 1 at run time proves the root is chosen dynamically. + const absl::string_view kModuleStr = R"( + HloModule test + apply_op { + x = u32[] parameter(0) + y = u32[] parameter(1) + ROOT apply_op = u32[] add(x, y) + } + ENTRY test_computation { + id = u32[] replica-id() + one = u32[] constant(1) + val = u32[] add(id, one) + p = u32[2] broadcast(val), dimensions={} + root = s32[1] constant({1}) + ROOT cr = u32[2] collective-reduce(p, root), replica_groups={{0, 1}}, + to_apply=apply_op, has_dynamic_root=true + } + )"; + const int64_t kNumReplicas = 2; + ASSERT_GE(device_count(), kNumReplicas) + << "Test requires at least " << kNumReplicas << " devices (" + << device_count() << " available)"; + + TF_ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + TF_ASSERT_OK_AND_ASSIGN(ExecutionResult execution_result, + ExecuteReplicated(std::move(module))); + + const HloModule* hlo_module = execution_result.optimized_module; + if (enable_async_) { + const HloInstruction* cr_start = + FindCollectiveStart(hlo_module, HloOpcode::kCollectiveReduce); + EXPECT_THAT(cr_start, NotNull()); + if (cr_start != nullptr) { + EXPECT_TRUE(Cast( + cr_start->async_wrapped_instruction()) + ->has_dynamic_root()); + } + } else { + const HloInstruction* cr_sync = + FindInstruction(hlo_module, HloOpcode::kCollectiveReduce); + EXPECT_THAT(cr_sync, NotNull()); + if (cr_sync != nullptr) { + EXPECT_TRUE( + Cast(cr_sync)->has_dynamic_root()); + } + } + + const std::vector& results = execution_result.results; + ASSERT_EQ(results.size(), kNumReplicas); + // Reduced sum (3) is written to the runtime-selected root rank (replica 1). + LiteralTestUtil::ExpectR1Equal({3, 3}, results[1]); +} + TEST_P(CollectivesModeOps, AllGather) { const absl::string_view kModuleStr = R"( HloModule test diff --git a/third_party/xla/xla/backends/gpu/tests/single_instruction.hlo b/third_party/xla/xla/backends/gpu/tests/single_instruction.hlo index 0e60c3a979a986..451c78fcd491d5 100644 --- a/third_party/xla/xla/backends/gpu/tests/single_instruction.hlo +++ b/third_party/xla/xla/backends/gpu/tests/single_instruction.hlo @@ -1,8 +1,10 @@ -// RUN: hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb --split-input-file | FileCheck %s -// RUN: hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s --check-prefixes=CHECK-SM80 -// RUN: hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s --check-prefixes=CHECK-SM90 +// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb --split-input-file | FileCheck %s %} +// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s --check-prefixes=CHECK-SM80 %} +// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s --check-prefixes=CHECK-SM90 %} +// RUN: %if IS_ROCM %{ hlo-opt %s --platform=gpu --stage=llvm --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb --split-input-file | FileCheck --check-prefixes=CHECK-GCN %s %} // CHECK-DAG: sqrt.approx.f32 +// CHECK-GCN: llvm.sqrt.f32 HloModule Test, is_scheduled=true @@ -19,6 +21,7 @@ ENTRY main { // ----- // CHECK-DAG: sqrt.approx.f32 +// CHECK-GCN: llvm.sqrt.f16 HloModule Test, is_scheduled=true @@ -35,6 +38,7 @@ ENTRY main { // ----- // CHECK-DAG: rsqrt.approx.f32 +// CHECK-GCN: llvm.amdgcn.rsq.f32 HloModule Test, is_scheduled=true @@ -52,6 +56,8 @@ ENTRY main { // CHECK: ex2.approx.ftz.f32 %[[APPROX:.*]], %{{.*}} // CHECK: mul.rn.f32 %{{.*}}, %[[APPROX]], %{{.*}} +// CHECK-GCN: llvm.exp.f32 +// CHECK-GCN-NOT: llvm.amdgcn.exp2 HloModule DoesntUseEx2ApproximationDirectly, is_scheduled=true @@ -70,6 +76,8 @@ ENTRY main { // CHECK-LABEL: exp_fusion( // CHECK: ex2.approx.f32 // CHECK-NOT: mul +// CHECK-GCN-LABEL: exp_fusion( +// CHECK-GCN: llvm.amdgcn.exp2.f32 HloModule UsesEx2ApproximationDirectly, is_scheduled=true @@ -87,6 +95,9 @@ ENTRY main { // CHECK-LABEL: log_fusion( // CHECK-NOT: lg2.approx.f32 +// CHECK-GCN-LABEL: log_fusion( +// CHECK-GCN: llvm.log.f32 +// CHECK-GCN-NOT: llvm.amdgcn.log HloModule DoesNotUseLogApproximation, is_scheduled=true @@ -104,6 +115,8 @@ ENTRY main { // CHECK-LABEL: log_fusion( // CHECK: lg2.approx.f32 +// CHECK-GCN-LABEL: log_fusion( +// CHECK-GCN: llvm.amdgcn.log.f32 HloModule UsesLogApproximation, is_scheduled=true @@ -120,6 +133,7 @@ ENTRY main { // ----- // CHECK-SM80: min.NaN.f32 +// CHECK-GCN: llvm.minimum.f32 HloModule Test, is_scheduled=true @@ -140,6 +154,7 @@ ENTRY main { // CHECK-SM80: cvt.rn.f32.s16 // CHECK-SM80: cvt.rn.bf16.f32 // CHECK-SM90: cvt.rn.bf16.s16 +// CHECK-GCN: sitofp i16 {{.*}} to bfloat HloModule Test, is_scheduled=true diff --git a/third_party/xla/xla/benchmarks/core/benchmark.py b/third_party/xla/xla/benchmarks/core/benchmark.py index 6febedad609405..7d6230384a221d 100644 --- a/third_party/xla/xla/benchmarks/core/benchmark.py +++ b/third_party/xla/xla/benchmarks/core/benchmark.py @@ -42,6 +42,7 @@ "f32": jnp.float32, "f8e4m3fn": jnp.float8_e4m3fn, "f8e5m2": jnp.float8_e5m2, + "f4e2m1fn": jnp.float4_e2m1fn, "s2": jnp.int2, "s4": jnp.int4, "s8": jnp.int8, diff --git a/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc b/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc index fce0866c1aff6c..f5ca754ce5f89c 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/expand_integer_power.cc @@ -46,6 +46,24 @@ mlir::LogicalResult ExpandIntegerPower(mlir::math::IPowIOp op, mlir::PatternRewriter& rewriter) { llvm::SmallVector result_types(op->getResultTypes()); llvm::SmallVector arg_types(op->getOperandTypes()); + // When xla.is_unsigned is set, specify unsigned element types in arg_types so + // mapMhloOpToStdScalarOp detects unsigned semantics and omits the negative + // exponent check. result_types and operand values remain signless because + // the generated standard arith ops (muli, shrui, etc.) and the original + // math.ipowi operate on signless integer types. + if (op->hasAttr("xla.is_unsigned")) { + for (mlir::Type& type : arg_types) { + mlir::Type elem_type = mlir::getElementTypeOrSelf(type); + auto unsigned_elem_type = mlir::IntegerType::get( + op.getContext(), elem_type.getIntOrFloatBitWidth(), + mlir::IntegerType::Unsigned); + if (auto shaped_type = mlir::dyn_cast(type)) { + type = shaped_type.clone(unsigned_elem_type); + } else { + type = unsigned_elem_type; + } + } + } mlir::Value result = mlir::mhlo::impl::mapMhloOpToStdScalarOp( op.getLoc(), result_types, arg_types, diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc index d167b17cd4c1dd..b64f626f12c522 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc @@ -492,8 +492,11 @@ struct RewriteTransferRead : OpRewritePattern { } }; -struct RewriteTensorInsert : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +class RewriteTensorInsert : public OpRewritePattern { + public: + RewriteTensorInsert(mlir::MLIRContext* context, const DeviceSpec& device_spec) + : OpRewritePattern(context), + device_spec_(device_spec) {} LogicalResult matchAndRewrite( mlir::tensor::InsertOp op, @@ -512,24 +515,11 @@ struct RewriteTensorInsert : OpRewritePattern { Type element_type = tensor_dest.getType().getElementType(); std::optional sub_byte_width = GetSubByteBitWidth(element_type); if (sub_byte_width) { - // We need to use directly op.getDest() as input, otherwise the following - // rewrite might remove the only user of it. - tensor_dest = op.getDest(); - // Value is_low_nibble; Value sub_byte_shift = nullptr; std::tie(linear_index, sub_byte_shift) = GetSubByteIndex(linear_index, *sub_byte_width, b); - // Technically we should half the number of elements when going to i8 - // element type, but it doesn't really matter because we only actually use - // the element type. Indexing is done by linear index, and GEP ops don't - // care about the number of elements. The tensor types will disappear - // completely after the LowerTensors pass. Type ty = b.getI8Type(); - Type tensor_ty = tensor_dest.getType().clone(ty); - auto tensor_dest_i8 = - UnrealizedConversionCastOp::create(b, tensor_ty, tensor_dest) - .getResult(0); if (scalar_value.getType() != b.getIntegerType(*sub_byte_width)) { scalar_value = arith::BitcastOp::create( b, b.getIntegerType(*sub_byte_width), scalar_value); @@ -548,22 +538,56 @@ struct RewriteTensorInsert : OpRewritePattern { Value shifted_value = mlir::arith::ShLIOp::create(b, scalar_value, sub_byte_shift); - // We need AtomicRMWOp because it can happen that different threads try to - // access the same memory location. - auto atomic_rmw = AtomicRMWOp::create(b, tensor_dest_i8, linear_index); - mlir::ImplicitLocOpBuilder body_builder(atomic_rmw.getLoc(), - atomic_rmw.getBodyBuilder()); - Value current_value = atomic_rmw.getCurrentValue(); - Value masked_out_current_value = mlir::arith::AndIOp::create( - body_builder, current_value, inverse_mask); - Value new_value = mlir::arith::OrIOp::create( - body_builder, masked_out_current_value, shifted_value); - scf::YieldOp::create(body_builder, new_value); - Value casted_result = - UnrealizedConversionCastOp::create(b, tensor_dest.getType(), - atomic_rmw.getResult()) - .getResult(0); - op.replaceAllUsesWith(casted_result); + if (device_spec_.IsCpu()) { + // On CPU, workgroup tiles are byte-aligned (see `GetWorkDimensions`), + // so distinct threads never access the same byte. + auto gep = CreateGep(tensor_dest, linear_index, b); + Value current_value = ml::LoadOp::create(b, ty, gep); + Value masked_out_current_value = + mlir::arith::AndIOp::create(b, current_value, inverse_mask); + Value new_value = mlir::arith::OrIOp::create( + b, masked_out_current_value, shifted_value); + auto store_op = ml::StoreOp::create(b, new_value, gep); + if (auto alias_scope_attr = op->getAttrOfType( + ml::LLVMDialect::getAliasScopesAttrName())) { + store_op.setAliasScopesAttr(alias_scope_attr); + } + if (auto no_alias_attr = op->getAttrOfType( + ml::LLVMDialect::getNoAliasAttrName())) { + store_op.setNoaliasScopesAttr(no_alias_attr); + } + op.replaceAllUsesWith(op.getDest()); + } else { + // We need to use directly op.getDest() as input, otherwise the + // following rewrite might remove the only user of it. + tensor_dest = op.getDest(); + // Technically we should halve the number of elements when going to i8 + // element type, but it doesn't really matter because we only actually + // use the element type. Indexing is done by linear index, and GEP ops + // don't care about the number of elements. The tensor types will + // disappear completely after the LowerTensors pass. + Type tensor_ty = tensor_dest.getType().clone(ty); + auto tensor_dest_i8 = + UnrealizedConversionCastOp::create(b, tensor_ty, tensor_dest) + .getResult(0); + + // We need AtomicRMWOp because it can happen that different threads try + // to access the same memory location. + auto atomic_rmw = AtomicRMWOp::create(b, tensor_dest_i8, linear_index); + mlir::ImplicitLocOpBuilder body_builder(atomic_rmw.getLoc(), + atomic_rmw.getBodyBuilder()); + Value current_value = atomic_rmw.getCurrentValue(); + Value masked_out_current_value = mlir::arith::AndIOp::create( + body_builder, current_value, inverse_mask); + Value new_value = mlir::arith::OrIOp::create( + body_builder, masked_out_current_value, shifted_value); + scf::YieldOp::create(body_builder, new_value); + Value casted_result = + UnrealizedConversionCastOp::create(b, tensor_dest.getType(), + atomic_rmw.getResult()) + .getResult(0); + op.replaceAllUsesWith(casted_result); + } } else { auto gep = CreateGep(tensor_dest, linear_index, b); mlir::LLVMTypeConverter converter(getContext()); @@ -588,6 +612,9 @@ struct RewriteTensorInsert : OpRewritePattern { op.erase(); return success(); } + + private: + const DeviceSpec& device_spec_; }; struct RewriteTransferWrite : OpRewritePattern { @@ -1230,12 +1257,18 @@ class RewriteAtomicRMW : public OpRewritePattern { // values from the memory. This can avoid out of bound memory accesses if // tensor buffers are 4 byte aligned and have a size of 4N, an assumption // that the runtime can guarantee. + // + // On CPUs, 8-bit and 16-bit atomic compare-and-swap operations are supported + // natively, and runtime buffers are not guaranteed to have a size of 4N. We + // therefore use a minimum atomic width of 8 bits on CPU and only widen + // sub-byte types to i8 without rounding the byte address down. void rewriteAsAtomicCAS(AtomicRMWOp op, mlir::PatternRewriter& rewriter) const { Location loc = op.getLoc(); auto input = op.getInput(); - // Use 32-bit atomic type for small input types. + // Use 32-bit atomic type for small input types on GPU, and 8-bit minimum + // atomic type on CPU. Type result_ty = op.getResult().getType().getElementType(); int result_size; if (auto complex_ty = mlir::dyn_cast(result_ty)) { @@ -1244,9 +1277,10 @@ class RewriteAtomicRMW : public OpRewritePattern { result_size = result_ty.getIntOrFloatBitWidth(); } - bool small_type = result_size < 32; - Type atomic_ty = - mlir::IntegerType::get(op.getContext(), small_type ? 32 : result_size); + int min_atomic_size = device_spec_.IsCpu() ? 8 : 32; + bool small_type = result_size < min_atomic_size; + Type atomic_ty = mlir::IntegerType::get( + op.getContext(), small_type ? min_atomic_size : result_size); // Calculate load address for the input. mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter); @@ -1261,29 +1295,37 @@ class RewriteAtomicRMW : public OpRewritePattern { Value addr = CreateGep(input, linear_index, b); Value shift, mask; if (small_type) { - // Update input pointer by discarding the last two bits - i.e. align to - // 32-bit boundary for small input types (will not result in OOB, as the - // input alignment is at least 32 bits). - Type addr_int_ty = rewriter.getI64Type(); - Value addr_int = ml::PtrToIntOp::create(rewriter, loc, addr_int_ty, addr); - Value addr_offset = ml::AndOp::create( - rewriter, loc, addr_int, - ml::ConstantOp::create(rewriter, loc, addr_int_ty, 3)); - Value index = ml::MulOp::create( - rewriter, loc, addr_offset, - ml::ConstantOp::create(rewriter, loc, addr_int_ty, -1)); - addr = - ml::GEPOp::create(rewriter, loc, addr.getType(), rewriter.getI8Type(), - addr, index, mlir::LLVM::GEPNoWrapFlags::inbounds); - - Value offset = ml::TruncOp::create(rewriter, loc, atomic_ty, addr_offset); - shift = ml::MulOp::create( - rewriter, loc, offset, - ml::ConstantOp::create(rewriter, loc, offset.getType(), 8)); - if (sub_byte_bit_width) { - sub_byte_shift = - ml::ZExtOp::create(rewriter, loc, shift.getType(), sub_byte_shift); - shift = ml::AddOp::create(rewriter, loc, shift, sub_byte_shift); + if (min_atomic_size > 8) { + // Update input pointer by discarding the last two bits - i.e. align to + // 32-bit boundary for small input types (will not result in OOB, as the + // input alignment is at least 32 bits). + Type addr_int_ty = rewriter.getI64Type(); + Value addr_int = + ml::PtrToIntOp::create(rewriter, loc, addr_int_ty, addr); + Value addr_offset = ml::AndOp::create( + rewriter, loc, addr_int, + ml::ConstantOp::create(rewriter, loc, addr_int_ty, 3)); + Value index = ml::MulOp::create( + rewriter, loc, addr_offset, + ml::ConstantOp::create(rewriter, loc, addr_int_ty, -1)); + addr = ml::GEPOp::create(rewriter, loc, addr.getType(), + rewriter.getI8Type(), addr, index, + mlir::LLVM::GEPNoWrapFlags::inbounds); + + Value offset = + ml::TruncOp::create(rewriter, loc, atomic_ty, addr_offset); + shift = ml::MulOp::create( + rewriter, loc, offset, + ml::ConstantOp::create(rewriter, loc, offset.getType(), 8)); + if (sub_byte_bit_width) { + sub_byte_shift = ml::ZExtOp::create(rewriter, loc, shift.getType(), + sub_byte_shift); + shift = ml::AddOp::create(rewriter, loc, shift, sub_byte_shift); + } + } else { + shift = sub_byte_bit_width + ? sub_byte_shift + : ml::ConstantOp::create(rewriter, loc, atomic_ty, 0); } // Compose the update mask. @@ -1436,12 +1478,13 @@ class LowerTensorsPass : public impl::LowerTensorsPassBase { MLIRContext* mlir_context = &getContext(); mlir::RewritePatternSet tensor_patterns(mlir_context); - tensor_patterns.add(mlir_context, device_spec_); - tensor_patterns.add( - mlir_context); + tensor_patterns.add(mlir_context, + device_spec_); + tensor_patterns + .add( + mlir_context); if (mlir::failed(mlir::applyPatternsGreedily(getOperation(), std::move(tensor_patterns)))) { signalPassFailure(); diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/expand_integer_power.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/expand_integer_power.mlir index df100c7358a4bc..5a4e6d4cfdcacd 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/tests/expand_integer_power.mlir +++ b/third_party/xla/xla/codegen/emitters/transforms/tests/expand_integer_power.mlir @@ -40,3 +40,29 @@ func.func @expand_integer_power_tensor( %0 = math.ipowi %arg0, %arg1 : tensor<6x4xi32> func.return %0 : tensor<6x4xi32> } + +// ----- + +// CHECK-LABEL: @expand_integer_power_unsigned +func.func @expand_integer_power_unsigned( + %arg0: vector<4xi32>, + %arg1: vector<4xi32>) -> vector<4xi32> { + // CHECK-NOT: math.ipowi + // CHECK: scf.for + // CHECK-NOT: arith.cmpi slt + %0 = math.ipowi %arg0, %arg1 {xla.is_unsigned} : vector<4xi32> + func.return %0 : vector<4xi32> +} + +// ----- + +// CHECK-LABEL: @expand_integer_power_unsigned_scalar +func.func @expand_integer_power_unsigned_scalar( + %arg0: i32, + %arg1: i32) -> i32 { + // CHECK-NOT: math.ipowi + // CHECK: scf.for + // CHECK-NOT: arith.cmpi slt + %0 = math.ipowi %arg0, %arg1 {xla.is_unsigned} : i32 + func.return %0 : i32 +} diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir index 5161313b86f202..97dcdd8c6bdf63 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir +++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir @@ -14,11 +14,11 @@ // ============================================================================== // RUN: emitters_opt %s --allow-unregistered-dialect -split-input-file \ // RUN: -xla-lower-tensors="gpu_device_info='cuda_compute_capability {major: 6}'" \ -// RUN: | FileCheck %s +// RUN: | FileCheck %s --check-prefixes=CHECK,CHECK-GPU // RUN: emitters_opt %s --allow-unregistered-dialect -split-input-file \ // RUN: -xla-lower-tensors="target_type=cpu" \ -// RUN: | FileCheck %s +// RUN: | FileCheck %s --check-prefixes=CHECK,CHECK-CPU // RUN: emitters_opt %s --allow-unregistered-dialect -split-input-file \ // RUN: -xla-lower-tensors="gpu_device_info='cuda_compute_capability {major: 6}'" \ @@ -287,21 +287,27 @@ func.func @atomic_rmw_f16(%in: tensor<8xf16>, %i: index) } // CHECK-LABEL: @atomic_rmw_f16 // CHECK: %[[ADDR:.*]] = llvm.getelementptr -// CHECK-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] -// CHECK-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} -// CHECK-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} -// CHECK-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]] -// CHECK: %[[INIT:.*]] = llvm.load %[[BASE]] -// CHECK-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) -// CHECK-NEXT: %[[VAR_SHIFT:.*]] = llvm.lshr %[[VAR]], %{{.*}} -// CHECK-NEXT: %[[VAR_TRUNC:.*]] = llvm.trunc %[[VAR_SHIFT]] -// CHECK-NEXT: arith.bitcast %[[VAR_TRUNC]] : i16 to f16 -// CHECK: %[[RES:.*]] = arith.bitcast %{{.*}} : f16 to i16 -// CHECK-NEXT: %[[RES_WIDE:.*]] = llvm.zext %[[RES]] -// CHECK-NEXT: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} -// CHECK-NEXT: %[[RES_SHIFT:.*]] = llvm.shl %[[RES_WIDE]], %{{.*}} -// CHECK-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] -// CHECK-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-GPU-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] +// CHECK-GPU-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} +// CHECK-GPU-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} +// CHECK-GPU-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]] +// CHECK-GPU: %[[INIT:.*]] = llvm.load %[[BASE]] +// CHECK-GPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) +// CHECK-GPU-NEXT: %[[VAR_SHIFT:.*]] = llvm.lshr %[[VAR]], %{{.*}} +// CHECK-GPU-NEXT: %[[VAR_TRUNC:.*]] = llvm.trunc %[[VAR_SHIFT]] +// CHECK-GPU-NEXT: arith.bitcast %[[VAR_TRUNC]] : i16 to f16 +// CHECK-GPU: %[[RES:.*]] = arith.bitcast %{{.*}} : f16 to i16 +// CHECK-GPU-NEXT: %[[RES_WIDE:.*]] = llvm.zext %[[RES]] +// CHECK-GPU-NEXT: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} +// CHECK-GPU-NEXT: %[[RES_SHIFT:.*]] = llvm.shl %[[RES_WIDE]], %{{.*}} +// CHECK-GPU-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] +// CHECK-GPU-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-CPU-NEXT: %[[INIT:.*]] = llvm.load %[[ADDR]] : !llvm.ptr -> i16 +// CHECK-CPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) : (i16) -> i16 +// CHECK-CPU-NEXT: %[[CUR:.*]] = arith.bitcast %[[VAR]] : i16 to f16 +// CHECK-CPU-NEXT: %[[ADD:.*]] = arith.addf %[[CUR]], %{{.*}} : f16 +// CHECK-CPU-NEXT: %[[RES:.*]] = arith.bitcast %[[ADD]] : f16 to i16 +// CHECK-CPU-NEXT: llvm.cmpxchg %[[ADDR]], %[[VAR]], %[[RES]] monotonic monotonic : !llvm.ptr, i16 // ----- @@ -317,21 +323,27 @@ func.func @atomic_rmw_f8E4M3(%in: tensor<8xf8E4M3>, %i: index) } // CHECK-LABEL: @atomic_rmw_f8E4M3 // CHECK: %[[ADDR:.*]] = llvm.getelementptr -// CHECK-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] -// CHECK-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} -// CHECK-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} -// CHECK-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]] -// CHECK: %[[INIT:.*]] = llvm.load %[[BASE]] -// CHECK-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) -// CHECK-NEXT: %[[VAR_SHIFT:.*]] = llvm.lshr %[[VAR]], %{{.*}} -// CHECK-NEXT: %[[VAR_TRUNC:.*]] = llvm.trunc %[[VAR_SHIFT]] -// CHECK-NEXT: arith.bitcast %[[VAR_TRUNC]] : i8 to f8E4M3 -// CHECK: %[[RES:.*]] = arith.bitcast %{{.*}} : f8E4M3 to i8 -// CHECK-NEXT: %[[RES_WIDE:.*]] = llvm.zext %[[RES]] -// CHECK-DAG: %[[RES_SHIFT:.*]] = llvm.shl %[[RES_WIDE]], %{{.*}} -// CHECK-DAG: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} -// CHECK-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] -// CHECK-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-GPU-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] +// CHECK-GPU-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} +// CHECK-GPU-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} +// CHECK-GPU-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]] +// CHECK-GPU: %[[INIT:.*]] = llvm.load %[[BASE]] +// CHECK-GPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) +// CHECK-GPU-NEXT: %[[VAR_SHIFT:.*]] = llvm.lshr %[[VAR]], %{{.*}} +// CHECK-GPU-NEXT: %[[VAR_TRUNC:.*]] = llvm.trunc %[[VAR_SHIFT]] +// CHECK-GPU-NEXT: arith.bitcast %[[VAR_TRUNC]] : i8 to f8E4M3 +// CHECK-GPU: %[[RES:.*]] = arith.bitcast %{{.*}} : f8E4M3 to i8 +// CHECK-GPU-NEXT: %[[RES_WIDE:.*]] = llvm.zext %[[RES]] +// CHECK-GPU-DAG: %[[RES_SHIFT:.*]] = llvm.shl %[[RES_WIDE]], %{{.*}} +// CHECK-GPU-DAG: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} +// CHECK-GPU-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] +// CHECK-GPU-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-CPU-NEXT: %[[INIT:.*]] = llvm.load %[[ADDR]] : !llvm.ptr -> i8 +// CHECK-CPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) : (i8) -> i8 +// CHECK-CPU-NEXT: %[[CUR:.*]] = arith.bitcast %[[VAR]] : i8 to f8E4M3 +// CHECK-CPU-NEXT: %[[ADD:.*]] = arith.addf %[[CUR]], %{{.*}} : f8E4M3 +// CHECK-CPU-NEXT: %[[RES:.*]] = arith.bitcast %[[ADD]] : f8E4M3 to i8 +// CHECK-CPU-NEXT: llvm.cmpxchg %[[ADDR]], %[[VAR]], %[[RES]] monotonic monotonic : !llvm.ptr, i8 // ----- @@ -347,45 +359,98 @@ func.func @atomic_rmw_i4(%in: tensor<8xi4>, %i: index) -> (tensor<8xi4>) { // CHECK-LABEL: @atomic_rmw_i4( // CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr, %[[ARG1:.*]]: index) { -// CHECK: %[[C1_I4:.*]] = arith.constant 1 : i4 -// CHECK: %[[C1_I64:.*]] = arith.constant 1 : i64 -// CHECK: %[[C4_I8:.*]] = arith.constant 4 : i8 -// CHECK: %[[C3_I64:.*]] = llvm.mlir.constant(3 : i64) : i64 -// CHECK: %[[C_NEG1_I64:.*]] = llvm.mlir.constant(-1 : i64) : i64 -// CHECK: %[[C8_I32:.*]] = llvm.mlir.constant(8 : i32) : i32 -// CHECK: %[[C_NEG1_I32:.*]] = llvm.mlir.constant(-1 : i32) : i32 -// CHECK: %[[C15_I32:.*]] = llvm.mlir.constant(15 : i32) : i32 -// CHECK: %[[C_TRUE:.*]] = llvm.mlir.constant(true) : i1 -// CHECK: %[[ARG1_I64:.*]] = arith.index_castui %[[ARG1]] : index to i64 -// CHECK: %[[AND_1:.*]] = arith.andi %[[ARG1_I64]], %[[C1_I64]] : i64 -// CHECK: %[[TRUNC_I8:.*]] = arith.trunci %[[AND_1]] : i64 to i8 -// CHECK: %[[MUL_4:.*]] = arith.muli %[[TRUNC_I8]], %[[C4_I8]] : i8 -// CHECK: %[[SHR_1:.*]] = arith.shrui %[[ARG1_I64]], %[[C1_I64]] : i64 -// CHECK: %[[GEP_BASE:.*]] = llvm.getelementptr inbounds %[[ARG0]][0, %[[SHR_1]]] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<4 x i8> -// CHECK: %[[PTR_INT:.*]] = llvm.ptrtoint %[[GEP_BASE]] : !llvm.ptr to i64 -// CHECK: %[[AND_3:.*]] = llvm.and %[[PTR_INT]], %[[C3_I64]] : i64 -// CHECK: %[[MUL_NEG1:.*]] = llvm.mul %[[AND_3]], %[[C_NEG1_I64]] : i64 -// CHECK: %[[GEP_OFFSET:.*]] = llvm.getelementptr inbounds %[[GEP_BASE]][%[[MUL_NEG1]]] : (!llvm.ptr, i64) -> !llvm.ptr, i8 -// CHECK: %[[TRUNC_I32:.*]] = llvm.trunc %[[AND_3]] : i64 to i32 -// CHECK: %[[MUL_8:.*]] = llvm.mul %[[TRUNC_I32]], %[[C8_I32]] : i32 -// CHECK: %[[ZEXT_I32:.*]] = llvm.zext %[[MUL_4]] : i8 to i32 -// CHECK: %[[ADD_19:.*]] = llvm.add %[[MUL_8]], %[[ZEXT_I32]] : i32 -// CHECK: %[[SHL_20:.*]] = llvm.shl %[[C15_I32]], %[[ADD_19]] : i32 -// CHECK: %[[XOR_21:.*]] = llvm.xor %[[C_NEG1_I32]], %[[SHL_20]] : i32 -// CHECK: %[[LOADED_I32:.*]] = llvm.load %[[GEP_OFFSET]] : !llvm.ptr -> i32 -// CHECK: %[[WHILE_RES:.*]] = scf.while (%[[ARG2:.*]] = %[[LOADED_I32]]) : (i32) -> i32 { -// CHECK: %[[LSHR:.*]] = llvm.lshr %[[ARG2]], %[[ADD_19]] : i32 -// CHECK: %[[TRUNC_I4:.*]] = llvm.trunc %[[LSHR]] : i32 to i4 -// CHECK: %[[ADDI:.*]] = arith.addi %[[TRUNC_I4]], %[[C1_I4]] : i4 -// CHECK: %[[ZEXT_I32_2:.*]] = llvm.zext %[[ADDI]] : i4 to i32 -// CHECK: %[[AND_MASK_I32:.*]] = llvm.and %[[ARG2]], %[[XOR_21]] : i32 -// CHECK: %[[SHL_NEW:.*]] = llvm.shl %[[ZEXT_I32_2]], %[[ADD_19]] : i32 -// CHECK: %[[OR_NEW:.*]] = llvm.or %[[AND_MASK_I32]], %[[SHL_NEW]] : i32 -// CHECK: %[[CMPXCHG:.*]] = llvm.cmpxchg %[[GEP_OFFSET]], %[[ARG2]], %[[OR_NEW]] monotonic monotonic : !llvm.ptr, i32 -// CHECK: %[[EXTRACT0:.*]] = llvm.extractvalue %[[CMPXCHG]][0] : !llvm.struct<(i32, i1)> -// CHECK: %[[EXTRACT1:.*]] = llvm.extractvalue %[[CMPXCHG]][1] : !llvm.struct<(i32, i1)> -// CHECK: %[[XOR_34:.*]] = llvm.xor %[[EXTRACT1]], %[[C_TRUE]] : i1 -// CHECK: scf.condition +// CHECK-GPU: %[[C1_I4:.*]] = arith.constant 1 : i4 +// CHECK-GPU: %[[C1_I64:.*]] = arith.constant 1 : i64 +// CHECK-GPU: %[[C4_I8:.*]] = arith.constant 4 : i8 +// CHECK-GPU: %[[C3_I64:.*]] = llvm.mlir.constant(3 : i64) : i64 +// CHECK-GPU: %[[C_NEG1_I64:.*]] = llvm.mlir.constant(-1 : i64) : i64 +// CHECK-GPU: %[[C8_I32:.*]] = llvm.mlir.constant(8 : i32) : i32 +// CHECK-GPU: %[[C_NEG1_I32:.*]] = llvm.mlir.constant(-1 : i32) : i32 +// CHECK-GPU: %[[C15_I32:.*]] = llvm.mlir.constant(15 : i32) : i32 +// CHECK-GPU: %[[C_TRUE:.*]] = llvm.mlir.constant(true) : i1 +// CHECK-GPU: %[[ARG1_I64:.*]] = arith.index_castui %[[ARG1]] : index to i64 +// CHECK-GPU: %[[AND_1:.*]] = arith.andi %[[ARG1_I64]], %[[C1_I64]] : i64 +// CHECK-GPU: %[[TRUNC_I8:.*]] = arith.trunci %[[AND_1]] : i64 to i8 +// CHECK-GPU: %[[MUL_4:.*]] = arith.muli %[[TRUNC_I8]], %[[C4_I8]] : i8 +// CHECK-GPU: %[[SHR_1:.*]] = arith.shrui %[[ARG1_I64]], %[[C1_I64]] : i64 +// CHECK-GPU: %[[GEP_BASE:.*]] = llvm.getelementptr inbounds %[[ARG0]][0, %[[SHR_1]]] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<4 x i8> +// CHECK-GPU: %[[PTR_INT:.*]] = llvm.ptrtoint %[[GEP_BASE]] : !llvm.ptr to i64 +// CHECK-GPU: %[[AND_3:.*]] = llvm.and %[[PTR_INT]], %[[C3_I64]] : i64 +// CHECK-GPU: %[[MUL_NEG1:.*]] = llvm.mul %[[AND_3]], %[[C_NEG1_I64]] : i64 +// CHECK-GPU: %[[GEP_OFFSET:.*]] = llvm.getelementptr inbounds %[[GEP_BASE]][%[[MUL_NEG1]]] : (!llvm.ptr, i64) -> !llvm.ptr, i8 +// CHECK-GPU: %[[TRUNC_I32:.*]] = llvm.trunc %[[AND_3]] : i64 to i32 +// CHECK-GPU: %[[MUL_8:.*]] = llvm.mul %[[TRUNC_I32]], %[[C8_I32]] : i32 +// CHECK-GPU: %[[ZEXT_I32:.*]] = llvm.zext %[[MUL_4]] : i8 to i32 +// CHECK-GPU: %[[ADD_19:.*]] = llvm.add %[[MUL_8]], %[[ZEXT_I32]] : i32 +// CHECK-GPU: %[[SHL_20:.*]] = llvm.shl %[[C15_I32]], %[[ADD_19]] : i32 +// CHECK-GPU: %[[XOR_21:.*]] = llvm.xor %[[C_NEG1_I32]], %[[SHL_20]] : i32 +// CHECK-GPU: %[[LOADED_I32:.*]] = llvm.load %[[GEP_OFFSET]] : !llvm.ptr -> i32 +// CHECK-GPU: %[[WHILE_RES:.*]] = scf.while (%[[ARG2:.*]] = %[[LOADED_I32]]) : (i32) -> i32 { +// CHECK-GPU: %[[LSHR:.*]] = llvm.lshr %[[ARG2]], %[[ADD_19]] : i32 +// CHECK-GPU: %[[TRUNC_I4:.*]] = llvm.trunc %[[LSHR]] : i32 to i4 +// CHECK-GPU: %[[ADDI:.*]] = arith.addi %[[TRUNC_I4]], %[[C1_I4]] : i4 +// CHECK-GPU: %[[ZEXT_I32_2:.*]] = llvm.zext %[[ADDI]] : i4 to i32 +// CHECK-GPU: %[[AND_MASK_I32:.*]] = llvm.and %[[ARG2]], %[[XOR_21]] : i32 +// CHECK-GPU: %[[SHL_NEW:.*]] = llvm.shl %[[ZEXT_I32_2]], %[[ADD_19]] : i32 +// CHECK-GPU: %[[OR_NEW:.*]] = llvm.or %[[AND_MASK_I32]], %[[SHL_NEW]] : i32 +// CHECK-GPU: %[[CMPXCHG:.*]] = llvm.cmpxchg %[[GEP_OFFSET]], %[[ARG2]], %[[OR_NEW]] monotonic monotonic : !llvm.ptr, i32 +// CHECK-GPU: %[[EXTRACT0:.*]] = llvm.extractvalue %[[CMPXCHG]][0] : !llvm.struct<(i32, i1)> +// CHECK-GPU: %[[EXTRACT1:.*]] = llvm.extractvalue %[[CMPXCHG]][1] : !llvm.struct<(i32, i1)> +// CHECK-GPU: %[[XOR_34:.*]] = llvm.xor %[[EXTRACT1]], %[[C_TRUE]] : i1 +// CHECK-GPU: scf.condition +// CHECK-CPU: %[[C1_I4:.*]] = arith.constant 1 : i4 +// CHECK-CPU: %[[C1_I64:.*]] = arith.constant 1 : i64 +// CHECK-CPU: %[[C4_I8:.*]] = arith.constant 4 : i8 +// CHECK-CPU: %[[C_NEG1_I8:.*]] = llvm.mlir.constant(-1 : i8) : i8 +// CHECK-CPU: %[[C15_I8:.*]] = llvm.mlir.constant(15 : i8) : i8 +// CHECK-CPU: %[[C_TRUE:.*]] = llvm.mlir.constant(true) : i1 +// CHECK-CPU: %[[ARG1_I64:.*]] = arith.index_castui %[[ARG1]] : index to i64 +// CHECK-CPU: %[[AND_1:.*]] = arith.andi %[[ARG1_I64]], %[[C1_I64]] : i64 +// CHECK-CPU: %[[TRUNC_I8:.*]] = arith.trunci %[[AND_1]] : i64 to i8 +// CHECK-CPU: %[[MUL_4:.*]] = arith.muli %[[TRUNC_I8]], %[[C4_I8]] : i8 +// CHECK-CPU: %[[SHR_1:.*]] = arith.shrui %[[ARG1_I64]], %[[C1_I64]] : i64 +// CHECK-CPU: %[[GEP_BASE:.*]] = llvm.getelementptr inbounds %[[ARG0]][0, %[[SHR_1]]] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<4 x i8> +// CHECK-CPU-NEXT: %[[SHL_MASK:.*]] = llvm.shl %[[C15_I8]], %[[MUL_4]] : i8 +// CHECK-CPU-NEXT: %[[INV_MASK:.*]] = llvm.xor %[[C_NEG1_I8]], %[[SHL_MASK]] : i8 +// CHECK-CPU-NEXT: %[[LOADED_I8:.*]] = llvm.load %[[GEP_BASE]] : !llvm.ptr -> i8 +// CHECK-CPU-NEXT: %[[WHILE_RES:.*]] = scf.while (%[[ARG2:.*]] = %[[LOADED_I8]]) : (i8) -> i8 { +// CHECK-CPU-NEXT: %[[LSHR:.*]] = llvm.lshr %[[ARG2]], %[[MUL_4]] : i8 +// CHECK-CPU-NEXT: %[[TRUNC_I4:.*]] = llvm.trunc %[[LSHR]] : i8 to i4 +// CHECK-CPU-NEXT: %[[ADDI:.*]] = arith.addi %[[TRUNC_I4]], %[[C1_I4]] : i4 +// CHECK-CPU-NEXT: %[[ZEXT_I8:.*]] = llvm.zext %[[ADDI]] : i4 to i8 +// CHECK-CPU-NEXT: %[[AND_MASK_I8:.*]] = llvm.and %[[ARG2]], %[[INV_MASK]] : i8 +// CHECK-CPU-NEXT: %[[SHL_NEW:.*]] = llvm.shl %[[ZEXT_I8]], %[[MUL_4]] : i8 +// CHECK-CPU-NEXT: %[[OR_NEW:.*]] = llvm.or %[[AND_MASK_I8]], %[[SHL_NEW]] : i8 +// CHECK-CPU-NEXT: %[[CMPXCHG:.*]] = llvm.cmpxchg %[[GEP_BASE]], %[[ARG2]], %[[OR_NEW]] monotonic monotonic : !llvm.ptr, i8 +// CHECK-CPU-NEXT: %[[EXTRACT0:.*]] = llvm.extractvalue %[[CMPXCHG]][0] : !llvm.struct<(i8, i1)> +// CHECK-CPU-NEXT: %[[EXTRACT1:.*]] = llvm.extractvalue %[[CMPXCHG]][1] : !llvm.struct<(i8, i1)> +// CHECK-CPU-NEXT: %[[XOR_OK:.*]] = llvm.xor %[[EXTRACT1]], %[[C_TRUE]] : i1 +// CHECK-CPU-NEXT: scf.condition + +// ----- + +func.func @atomic_rmw_i1(%in: tensor<8xi1>, %i: index) -> (tensor<8xi1>) { + %ret = xla.atomic_rmw %in[%i] : tensor<8xi1> { + ^bb0(%current : i1): + %c1 = arith.constant true + %xor = arith.xori %current, %c1 : i1 + xla.yield %xor : i1 + } + return %ret : tensor<8xi1> +} +// CHECK-LABEL: @atomic_rmw_i1 +// CHECK: %[[ADDR:.*]] = llvm.getelementptr inbounds %{{.*}}[0, %{{.*}}] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<8 x i1> +// CHECK-GPU-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] +// CHECK-GPU-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} +// CHECK-GPU-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} +// CHECK-GPU-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]][%[[INDEX]]] +// CHECK-GPU: %[[INIT:.*]] = llvm.load %[[BASE]] : !llvm.ptr -> i32 +// CHECK-GPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) : (i32) -> i32 +// CHECK-GPU: llvm.cmpxchg %[[BASE]], %[[VAR]], %{{.*}} monotonic monotonic : !llvm.ptr, i32 +// CHECK-CPU-NOT: llvm.ptrtoint +// CHECK-CPU: %[[INIT:.*]] = llvm.load %[[ADDR]] : !llvm.ptr -> i8 +// CHECK-CPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) : (i8) -> i8 +// CHECK-CPU: llvm.cmpxchg %[[ADDR]], %[[VAR]], %{{.*}} monotonic monotonic : !llvm.ptr, i8 // ----- @@ -400,16 +465,19 @@ func.func @atomic_rmw_overwrite(%in: tensor<8xf16>, %i: index) } // CHECK-LABEL: @atomic_rmw_overwrite // CHECK: %[[ADDR:.*]] = llvm.getelementptr -// CHECK-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] -// CHECK-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} -// CHECK-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} -// CHECK-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]][%[[INDEX]]] -// CHECK: %[[INIT:.*]] = llvm.load %[[BASE]] -// CHECK-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) -// CHECK-DAG: %[[RES_SHIFT:.*]] = llvm.shl %{{.*}}, %{{.*}} -// CHECK-DAG: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} -// CHECK-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] -// CHECK-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-GPU-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] +// CHECK-GPU-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} +// CHECK-GPU-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} +// CHECK-GPU-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]][%[[INDEX]]] +// CHECK-GPU: %[[INIT:.*]] = llvm.load %[[BASE]] +// CHECK-GPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) +// CHECK-GPU-DAG: %[[RES_SHIFT:.*]] = llvm.shl %{{.*}}, %{{.*}} +// CHECK-GPU-DAG: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} +// CHECK-GPU-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] +// CHECK-GPU-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-CPU-NEXT: %[[INIT:.*]] = llvm.load %[[ADDR]] : !llvm.ptr -> i16 +// CHECK-CPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) : (i16) -> i16 +// CHECK-CPU-NEXT: llvm.cmpxchg %[[ADDR]], %[[VAR]], %{{.*}} monotonic monotonic : !llvm.ptr, i16 // ----- @@ -424,16 +492,19 @@ func.func @atomic_rmw_overwrite_f8E4M3(%in: tensor<8xf8E4M3>, %i: index) } // CHECK-LABEL: @atomic_rmw_overwrite_f8E4M3 // CHECK: %[[ADDR:.*]] = llvm.getelementptr -// CHECK-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] -// CHECK-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} -// CHECK-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} -// CHECK-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]][%[[INDEX]]] -// CHECK: %[[INIT:.*]] = llvm.load %[[BASE]] -// CHECK-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) -// CHECK-DAG: %[[RES_SHIFT:.*]] = llvm.shl %{{.*}}, %{{.*}} -// CHECK-DAG: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} -// CHECK-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] -// CHECK-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-GPU-NEXT: %[[ADDR_INT:.*]] = llvm.ptrtoint %[[ADDR]] +// CHECK-GPU-NEXT: %[[OFFSET:.*]] = llvm.and %[[ADDR_INT]], %{{.*}} +// CHECK-GPU-NEXT: %[[INDEX:.*]] = llvm.mul %[[OFFSET]], %{{.*}} +// CHECK-GPU-NEXT: %[[BASE:.*]] = llvm.getelementptr inbounds %[[ADDR]][%[[INDEX]]] +// CHECK-GPU: %[[INIT:.*]] = llvm.load %[[BASE]] +// CHECK-GPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) +// CHECK-GPU-DAG: %[[RES_SHIFT:.*]] = llvm.shl %{{.*}}, %{{.*}} +// CHECK-GPU-DAG: %[[NEW_MASKED:.*]] = llvm.and %[[VAR]], %{{.*}} +// CHECK-GPU-NEXT: %[[NEW:.*]] = llvm.or %[[NEW_MASKED]], %[[RES_SHIFT]] +// CHECK-GPU-NEXT: llvm.cmpxchg %[[BASE]], %[[VAR]], %[[NEW]] +// CHECK-CPU-NEXT: %[[INIT:.*]] = llvm.load %[[ADDR]] : !llvm.ptr -> i8 +// CHECK-CPU-NEXT: scf.while (%[[VAR:.*]] = %[[INIT]]) : (i8) -> i8 +// CHECK-CPU-NEXT: llvm.cmpxchg %[[ADDR]], %[[VAR]], %{{.*}} monotonic monotonic : !llvm.ptr, i8 // ----- @@ -490,45 +561,52 @@ func.func @i4_store(%arg: tensor<10xi4>, %v: i4, %i: index) // CHECK-DAG: %[[C4_I8:.*]] = arith.constant 4 : i8 // CHECK-DAG: %[[C15_I8:.*]] = arith.constant 15 : i8 // CHECK-DAG: %[[C_NEG1_I8:.*]] = arith.constant -1 : i8 -// CHECK-DAG: %[[C3_I64:.*]] = llvm.mlir.constant(3 : i64) : i64 -// CHECK-DAG: %[[C_NEG1_I64:.*]] = llvm.mlir.constant(-1 : i64) : i64 -// CHECK-DAG: %[[C8_I32:.*]] = llvm.mlir.constant(8 : i32) : i32 -// CHECK-DAG: %[[C_NEG1_I32:.*]] = llvm.mlir.constant(-1 : i32) : i32 -// CHECK-DAG: %[[C255_I32:.*]] = llvm.mlir.constant(255 : i32) : i32 -// CHECK-DAG: %[[C_TRUE:.*]] = llvm.mlir.constant(true) : i1 -// CHECK: %[[CAST_ARG0:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : !llvm.ptr to tensor<10xi4> +// CHECK-GPU-DAG: %[[C3_I64:.*]] = llvm.mlir.constant(3 : i64) : i64 +// CHECK-GPU-DAG: %[[C_NEG1_I64:.*]] = llvm.mlir.constant(-1 : i64) : i64 +// CHECK-GPU-DAG: %[[C8_I32:.*]] = llvm.mlir.constant(8 : i32) : i32 +// CHECK-GPU-DAG: %[[C_NEG1_I32:.*]] = llvm.mlir.constant(-1 : i32) : i32 +// CHECK-GPU-DAG: %[[C255_I32:.*]] = llvm.mlir.constant(255 : i32) : i32 +// CHECK-GPU-DAG: %[[C_TRUE:.*]] = llvm.mlir.constant(true) : i1 +// CHECK-GPU: %[[CAST_ARG0:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : !llvm.ptr to tensor<10xi4> // CHECK: %[[ARG2_I64:.*]] = arith.index_castui %[[ARG2]] : index to i64 // CHECK: %[[AND_1:.*]] = arith.andi %[[ARG2_I64]], %[[C1_I64]] : i64 // CHECK: %[[TRUNC_I8:.*]] = arith.trunci %[[AND_1]] : i64 to i8 // CHECK: %[[MUL_4:.*]] = arith.muli %[[TRUNC_I8]], %[[C4_I8]] : i8 // CHECK: %[[SHR_1:.*]] = arith.shrui %[[ARG2_I64]], %[[C1_I64]] : i64 -// CHECK: %[[CAST_TENSOR:.*]] = builtin.unrealized_conversion_cast %[[CAST_ARG0]] : tensor<10xi4> to tensor<10xi8> // CHECK: %[[EXT_ARG1:.*]] = arith.extui %[[ARG1]] : i4 to i8 // CHECK: %[[SHL_15:.*]] = arith.shli %[[C15_I8]], %[[MUL_4]] : i8 // CHECK: %[[XOR_NEG1_I8:.*]] = arith.xori %[[SHL_15]], %[[C_NEG1_I8]] : i8 // CHECK: %[[SHL_EXT:.*]] = arith.shli %[[EXT_ARG1]], %[[MUL_4]] : i8 -// CHECK: %[[CAST_TENSOR_PTR:.*]] = builtin.unrealized_conversion_cast %[[CAST_TENSOR]] : tensor<10xi8> to !llvm.ptr -// CHECK: %[[GEP_BASE:.*]] = llvm.getelementptr inbounds %[[CAST_TENSOR_PTR]][0, %[[SHR_1]]] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<10 x i8> -// CHECK: %[[PTR_INT:.*]] = llvm.ptrtoint %[[GEP_BASE]] : !llvm.ptr to i64 -// CHECK: %[[AND_3:.*]] = llvm.and %[[PTR_INT]], %[[C3_I64]] : i64 -// CHECK: %[[MUL_NEG1:.*]] = llvm.mul %[[AND_3]], %[[C_NEG1_I64]] : i64 -// CHECK: %[[GEP_OFFSET:.*]] = llvm.getelementptr inbounds %[[GEP_BASE]][%[[MUL_NEG1]]] : (!llvm.ptr, i64) -> !llvm.ptr, i8 -// CHECK: %[[TRUNC_I32:.*]] = llvm.trunc %[[AND_3]] : i64 to i32 -// CHECK: %[[MUL_8:.*]] = llvm.mul %[[TRUNC_I32]], %[[C8_I32]] : i32 -// CHECK: %[[SHL_255:.*]] = llvm.shl %[[C255_I32]], %[[MUL_8]] : i32 -// CHECK: %[[XOR_NEG1_I32:.*]] = llvm.xor %[[C_NEG1_I32]], %[[SHL_255]] : i32 -// CHECK: %[[LOADED_I32:.*]] = llvm.load %[[GEP_OFFSET]] : !llvm.ptr -> i32 -// CHECK: %[[WHILE_RES:.*]] = scf.while (%[[INIT:.*]] = %[[LOADED_I32]]) : (i32) -> i32 { -// CHECK: %[[LSHR:.*]] = llvm.lshr %[[INIT]], %[[MUL_8]] : i32 -// CHECK: %[[TRUNC_I8_2:.*]] = llvm.trunc %[[LSHR]] : i32 to i8 -// CHECK: %[[AND_MASK_I8:.*]] = arith.andi %[[TRUNC_I8_2]], %[[XOR_NEG1_I8]] : i8 -// CHECK: %[[ORI_NEW_I8:.*]] = arith.ori %[[AND_MASK_I8]], %[[SHL_EXT]] : i8 -// CHECK: %[[ZEXT_I32:.*]] = llvm.zext %[[ORI_NEW_I8]] : i8 to i32 -// CHECK: %[[AND_MASK_I32:.*]] = llvm.and %[[INIT]], %[[XOR_NEG1_I32]] : i32 -// CHECK: %[[SHL_NEW:.*]] = llvm.shl %[[ZEXT_I32]], %[[MUL_8]] : i32 -// CHECK: %[[OR_NEW:.*]] = llvm.or %[[AND_MASK_I32]], %[[SHL_NEW]] : i32 -// CHECK: llvm.cmpxchg %[[GEP_OFFSET]], %[[INIT]], %[[OR_NEW]] monotonic monotonic : !llvm.ptr, i32 -// CHECK: scf.condition +// CHECK-GPU: %[[CAST_TENSOR:.*]] = builtin.unrealized_conversion_cast %[[CAST_ARG0]] : tensor<10xi4> to tensor<10xi8> +// CHECK-GPU: %[[CAST_TENSOR_PTR:.*]] = builtin.unrealized_conversion_cast %[[CAST_TENSOR]] : tensor<10xi8> to !llvm.ptr +// CHECK-GPU: %[[GEP_BASE:.*]] = llvm.getelementptr inbounds %[[CAST_TENSOR_PTR]][0, %[[SHR_1]]] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<10 x i8> +// CHECK-GPU: %[[PTR_INT:.*]] = llvm.ptrtoint %[[GEP_BASE]] : !llvm.ptr to i64 +// CHECK-GPU: %[[AND_3:.*]] = llvm.and %[[PTR_INT]], %[[C3_I64]] : i64 +// CHECK-GPU: %[[MUL_NEG1:.*]] = llvm.mul %[[AND_3]], %[[C_NEG1_I64]] : i64 +// CHECK-GPU: %[[GEP_OFFSET:.*]] = llvm.getelementptr inbounds %[[GEP_BASE]][%[[MUL_NEG1]]] : (!llvm.ptr, i64) -> !llvm.ptr, i8 +// CHECK-GPU: %[[TRUNC_I32:.*]] = llvm.trunc %[[AND_3]] : i64 to i32 +// CHECK-GPU: %[[MUL_8:.*]] = llvm.mul %[[TRUNC_I32]], %[[C8_I32]] : i32 +// CHECK-GPU: %[[SHL_255:.*]] = llvm.shl %[[C255_I32]], %[[MUL_8]] : i32 +// CHECK-GPU: %[[XOR_NEG1_I32:.*]] = llvm.xor %[[C_NEG1_I32]], %[[SHL_255]] : i32 +// CHECK-GPU: %[[LOADED_I32:.*]] = llvm.load %[[GEP_OFFSET]] : !llvm.ptr -> i32 +// CHECK-GPU: %[[WHILE_RES:.*]] = scf.while (%[[INIT:.*]] = %[[LOADED_I32]]) : (i32) -> i32 { +// CHECK-GPU: %[[LSHR:.*]] = llvm.lshr %[[INIT]], %[[MUL_8]] : i32 +// CHECK-GPU: %[[TRUNC_I8_2:.*]] = llvm.trunc %[[LSHR]] : i32 to i8 +// CHECK-GPU: %[[AND_MASK_I8:.*]] = arith.andi %[[TRUNC_I8_2]], %[[XOR_NEG1_I8]] : i8 +// CHECK-GPU: %[[ORI_NEW_I8:.*]] = arith.ori %[[AND_MASK_I8]], %[[SHL_EXT]] : i8 +// CHECK-GPU: %[[ZEXT_I32:.*]] = llvm.zext %[[ORI_NEW_I8]] : i8 to i32 +// CHECK-GPU: %[[AND_MASK_I32:.*]] = llvm.and %[[INIT]], %[[XOR_NEG1_I32]] : i32 +// CHECK-GPU: %[[SHL_NEW:.*]] = llvm.shl %[[ZEXT_I32]], %[[MUL_8]] : i32 +// CHECK-GPU: %[[OR_NEW:.*]] = llvm.or %[[AND_MASK_I32]], %[[SHL_NEW]] : i32 +// CHECK-GPU: llvm.cmpxchg %[[GEP_OFFSET]], %[[INIT]], %[[OR_NEW]] monotonic monotonic : !llvm.ptr, i32 +// CHECK-GPU: scf.condition +// CHECK-CPU-NOT: llvm.cmpxchg +// CHECK-CPU: %[[GEP_BASE:.*]] = llvm.getelementptr inbounds %[[ARG0]][0, %[[SHR_1]]] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<5 x i8> +// CHECK-CPU-NEXT: %[[LOADED_I8:.*]] = llvm.load %[[GEP_BASE]] : !llvm.ptr -> i8 +// CHECK-CPU-NEXT: %[[AND_MASK_I8:.*]] = arith.andi %[[LOADED_I8]], %[[XOR_NEG1_I8]] : i8 +// CHECK-CPU-NEXT: %[[ORI_NEW_I8:.*]] = arith.ori %[[AND_MASK_I8]], %[[SHL_EXT]] : i8 +// CHECK-CPU-NEXT: llvm.store %[[ORI_NEW_I8]], %[[GEP_BASE]] : i8, !llvm.ptr +// CHECK-CPU-NEXT: return // ----- diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h index 744216b987e032..5db376beaca61c 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h +++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h @@ -24,7 +24,9 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/types/span.h" #include "llvm/ADT/SmallVector.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -63,6 +65,33 @@ namespace xla::xtile { using TensorValue = mlir::TypedValue; static constexpr auto kTritonDivisibilityAttr = "tt.divisibility"; +// Maps sequential dimensions to their current value and its range. +using SequentialDimValueMap = + absl::flat_hash_map>; + +// Removes `dim_id` from `map` when destroyed. Mappings are scoped because the +// mapped values, e.g. loop induction variables, do not dominate code emitted +// outside of their loop. +class ScopedSequentialDimBinding { + public: + ScopedSequentialDimBinding(SequentialDimValueMap* map, + gpu::experimental::TiledDimId dim_id) + : map_(map), dim_id_(dim_id) {} + ScopedSequentialDimBinding(ScopedSequentialDimBinding&& other) noexcept + : map_(std::exchange(other.map_, nullptr)), dim_id_(other.dim_id_) {} + ScopedSequentialDimBinding& operator=(ScopedSequentialDimBinding&&) = delete; + ~ScopedSequentialDimBinding() { + if (map_ != nullptr) { + map_->erase(dim_id_); + } + } + + private: + SequentialDimValueMap* map_; + gpu::experimental::TiledDimId dim_id_; +}; + // Convenience class for holding the emitted values. class EmitterContext { public: @@ -108,13 +137,20 @@ class EmitterContext { return it->second; } - bool MapSymbolIdToSequentialDimValue( + // Maps `sequential_dim_id` to `value` until the returned object is destroyed. + // Returns an error if the dimension is already mapped. + absl::StatusOr + MapSymbolIdToSequentialDimValueScoped( gpu::experimental::TiledDimId sequential_dim_id, mlir::Value value, Interval interval) { - return sequential_dim_id_to_value_ - .insert( - std::make_pair(sequential_dim_id, std::make_pair(value, interval))) - .second; + if (!sequential_dim_id_to_value_ + .try_emplace(sequential_dim_id, value, interval) + .second) { + return absl::InternalError(absl::StrCat( + "Sequential dimension ", sequential_dim_id, " is already bound.")); + } + return ScopedSequentialDimBinding(&sequential_dim_id_to_value_, + sequential_dim_id); } // Evaluates tiling parameters for the given affine expressions, e.g. offsets. @@ -132,9 +168,7 @@ class EmitterContext { const HloFusionInstruction* fusion_ = nullptr; xtile::EntryFuncOp entry_func_; const gpu::experimental::TiledHloComputation& tiled_computation_; - absl::flat_hash_map> - sequential_dim_id_to_value_; + SequentialDimValueMap sequential_dim_id_to_value_; }; // Constructs and holds information needed to construct a tile. This information diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc index 0c7f9cfd2256d0..f3494bbe77856e 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc +++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc @@ -442,8 +442,10 @@ absl::StatusOr EmitDot(EmitterContext& emitter_ctx, const ge::TilingSpace::DimensionInfo& dim_info = tiled_dot.tile().tiling_space().GetDimensionInfo( *tiled_dot.hlo(), sequential_dim_ids.front()); - TF_RET_CHECK(emitter_ctx.MapSymbolIdToSequentialDimValue( - dim_info.id, iv, Interval{0, loop_iteration_count.front() - 1})); + ABSL_ASSIGN_OR_RETURN( + ScopedSequentialDimBinding iv_binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + dim_info.id, iv, Interval{0, loop_iteration_count.front() - 1})); // Emit the dot region. const ge::TiledHloInstruction* lhs_operand = tiled_dot.operand(0); @@ -537,8 +539,10 @@ absl::StatusOr EmitScaledDot( const ge::TilingSpace::DimensionInfo& dim_info = tiled_scaled_dot.tile().tiling_space().GetDimensionInfo( *tiled_scaled_dot.hlo(), sequential_dim_ids.front()); - TF_RET_CHECK(emitter_ctx.MapSymbolIdToSequentialDimValue( - dim_info.id, iv, Interval{0, loop_iteration_counts.front() - 1})); + ABSL_ASSIGN_OR_RETURN( + ScopedSequentialDimBinding iv_binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + dim_info.id, iv, Interval{0, loop_iteration_counts.front() - 1})); // Emit the dot region. const ge::TiledHloRegion& region = tiled_scaled_dot.hlo_regions().front(); @@ -751,8 +755,9 @@ absl::StatusOr EmitRaggedDot( // Register G sequential dim so EvaluateTilingParameters can resolve // RTVar offsets that depend on the G loop IV. - CHECK(emitter_ctx.MapSymbolIdToSequentialDimValue(g_dim_info.id, g_iv, - Interval{0, G - 1})); + ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding g_iv_binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + g_dim_info.id, g_iv, Interval{0, G - 1})); // Emit group_sizes tile (G-scoped). // Uses emit_gs which handles any HLO form: parameter, inlined constant, @@ -808,8 +813,9 @@ absl::StatusOr EmitRaggedDot( mlir::cast(k_for_op.getRegionIterArgs()[0]); // Register K sequential dim IV. - CHECK(emitter_ctx.MapSymbolIdToSequentialDimValue( - k_dim_info.id, k_iv, Interval{0, K_tiles - 1})); + ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding k_iv_binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + k_dim_info.id, k_iv, Interval{0, K_tiles - 1})); // Emit LHS and RHS tiles (K-scoped). // emit_operand handles the full instruction chain (parameter, @@ -1482,8 +1488,9 @@ absl::StatusOr> EmitScan( Value max_tile = MakeIndex(b, loop_count - 1); tile_idx = arith::SubIOp::create(b, max_tile, iv); } - emitter_ctx.MapSymbolIdToSequentialDimValue(dim_info.id, tile_idx, - Interval{0, loop_count - 1}); + ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding iv_binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + dim_info.id, tile_idx, Interval{0, loop_count - 1})); const auto& input_region = tiled_hlo_scan.hlo_regions().front(); ABSL_ASSIGN_OR_RETURN( @@ -1714,6 +1721,8 @@ absl::StatusOr EmitReduceWithRegion( auto body = [&]() -> absl::StatusOr { mlir::ImplicitLocOpBuilder nested_b(loc, loop_builder); + std::vector iv_bindings; + iv_bindings.reserve(sequential_dim_ids.size()); for (int i = 0; i < sequential_dim_ids.size(); ++i) { // Fully tiled dimensions were already bound to 0 by // EmitFullyTiledSequentialDimensions. Their single-trip loop is @@ -1724,10 +1733,11 @@ absl::StatusOr EmitReduceWithRegion( const ge::TilingSpace::DimensionInfo& dim_info = tiled_hlo.tile().tiling_space().GetDimensionInfo( *tiled_hlo.hlo(), sequential_dim_ids[i]); - TF_RET_CHECK(emitter_ctx.MapSymbolIdToSequentialDimValue( - dim_info.id, ivs[i], Interval{0, loop_iteration_counts[i] - 1})) - << "MapSymbolIdToSequentialDimValue failed for symbolic " - "dimension"; + ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding iv_binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + dim_info.id, ivs[i], + Interval{0, loop_iteration_counts[i] - 1})); + iv_bindings.push_back(std::move(iv_binding)); } ABSL_ASSIGN_OR_RETURN( @@ -1979,10 +1989,12 @@ absl::StatusOr> EmitTiledComputation( // [p, c, reduce], where p has a symbol dimension that is created by reduce. // To emit p we have to have a value for the symbol dimension. // Thus we emit sequential dimensions at the start as we know they will be -// trivially 0. -void EmitFullyTiledSequentialDimensions( +// trivially 0. The returned bindings must outlive the emission of the fusion. +absl::StatusOr> +EmitFullyTiledSequentialDimensions( ImplicitLocOpBuilder& b, EmitterContext& emitter_ctx, const ge::TiledHloComputation& tiled_computation) { + std::vector bindings; const auto& tiling_space = tiled_computation.tiling_space(); for (const auto& [dim_id, dim_info] : llvm::enumerate(tiling_space.dimensions())) { @@ -2001,10 +2013,14 @@ void EmitFullyTiledSequentialDimensions( << dim_info.dimension_size << " with tile size " << *dim_info.tile_size << " for hlo " << dim_info.hlo->name() << " to a new value 0"; - emitter_ctx.MapSymbolIdToSequentialDimValue( - ge::TiledDimId(dim_id), MakeIndex(b, 0), Interval{0, 0}); + ABSL_ASSIGN_OR_RETURN( + ScopedSequentialDimBinding binding, + emitter_ctx.MapSymbolIdToSequentialDimValueScoped( + ge::TiledDimId(dim_id), MakeIndex(b, 0), Interval{0, 0})); + bindings.push_back(std::move(binding)); } } + return bindings; } // Applies L2 tile reordering to the flat tile_id for kRaggedNonContracting and @@ -2240,7 +2256,9 @@ absl::Status EmitGeneric(ImplicitLocOpBuilder& b, schedule, fn, tiled_computation}; VLOG(2) << "EmitTiledComputation: " << tiled_computation.ToString(); - EmitFullyTiledSequentialDimensions(b, emitter_ctx, tiled_computation); + ABSL_ASSIGN_OR_RETURN( + std::vector fully_tiled_dim_bindings, + EmitFullyTiledSequentialDimensions(b, emitter_ctx, tiled_computation)); ABSL_ASSIGN_OR_RETURN( auto results, EmitTiledComputation(emitter_ctx, tiled_computation.tiled_root_region(), diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/lower_stablehlo_to_arith.cc b/third_party/xla/xla/codegen/xtile/ir/transforms/lower_stablehlo_to_arith.cc index ffa14465bec81b..d0735882dc00da 100644 --- a/third_party/xla/xla/codegen/xtile/ir/transforms/lower_stablehlo_to_arith.cc +++ b/third_party/xla/xla/codegen/xtile/ir/transforms/lower_stablehlo_to_arith.cc @@ -127,6 +127,9 @@ class LowerStableHloOpToArith : public mlir::OpRewritePattern { new_op = UnsignedIntArithOp::create(rewriter, op.getLoc(), signless_operands.front().getType(), signless_operands); + if constexpr (std::is_same_v) { + new_op->setAttr("xla.is_unsigned", rewriter.getUnitAttr()); + } rewriter.replaceOpWithNewOp( op, op.getResult().getType(), new_op->getResult(0)); diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/lower_stablehlo_to_arith.mlir b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/lower_stablehlo_to_arith.mlir index 4aea78d4bf62c5..327a8e2b910f2b 100644 --- a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/lower_stablehlo_to_arith.mlir +++ b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/lower_stablehlo_to_arith.mlir @@ -234,9 +234,19 @@ func.func @lower_pow_with_float_operands(%arg0 : tensor<2x4xf32>, %arg1 : tensor func.func @lower_pow_with_signless_operands(%arg0 : tensor<2x4xi32>, %arg1 : tensor<2x4xi32>) -> tensor<2x4xi32> { // CHECK: math.ipowi + // CHECK-NOT: xla.is_unsigned %0 = stablehlo.power %arg0, %arg1 : tensor<2x4xi32> return %0 : tensor<2x4xi32> } + +func.func @lower_pow_with_unsigned_operands(%arg0 : tensor<2x4xui32>, %arg1 : tensor<2x4xui32>) -> tensor<2x4xui32> { + // CHECK: builtin.unrealized_conversion_cast %{{.*}} : tensor<2x4xui32> to tensor<2x4xi32> + // CHECK: builtin.unrealized_conversion_cast %{{.*}} : tensor<2x4xui32> to tensor<2x4xi32> + // CHECK: math.ipowi {{.*}} {xla.is_unsigned} + // CHECK: builtin.unrealized_conversion_cast %{{.*}} : tensor<2x4xi32> to tensor<2x4xui32> + %0 = stablehlo.power %arg0, %arg1 : tensor<2x4xui32> + return %0 : tensor<2x4xui32> +} // ----- func.func @lower_convert_bf16_to_f32(%arg0: tensor<2x4xbf16>) -> tensor<2x4xf32> { diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index f9f29c90840802..4a5a969daac4eb 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -361,6 +361,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_enable_nccl_user_buffers_in_default_space(false); opts.set_xla_gpu_enable_allocator_spatial_partitioning(true); opts.set_xla_gpu_experimental_enable_nccl_symmetric_buffers(false); + opts.set_xla_gpu_experimental_emit_collective_reduce(false); opts.set_xla_gpu_enable_nccl_comm_splitting(true); opts.set_xla_gpu_nccl_init_max_rank_per_root_ratio(0); @@ -2401,6 +2402,15 @@ void MakeDebugOptionsFlags(std::vector* flag_list, set_xla_gpu_experimental_enable_nccl_symmetric_buffers), debug_options->xla_gpu_experimental_enable_nccl_symmetric_buffers(), "Enables NCCL symmetric buffer registration.")); + flag_list->push_back(tsl::Flag( + "xla_gpu_experimental_emit_collective_reduce", + bool_setter_for( + &DebugOptions::set_xla_gpu_experimental_emit_collective_reduce), + debug_options->xla_gpu_experimental_emit_collective_reduce(), + "Enables emitting a CollectiveReduceThunk for kCollectiveReduce HLO " + "instructions. Kept off by default to preserve the forward " + "compatibility window until the runtime support for the thunk has " + "rolled out.")); flag_list->push_back(tsl::Flag( "xla_enable_nccl_symmetric_buffers_for_collectives", setter_for_xla_enable_nccl_symmetric_buffers_for_collectives, diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index 37d0071dbbd980..3d353466b76f72 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -2731,15 +2731,15 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT if (!window) { window.emplace(); } + if (operands.empty()) { + TokenError("reduce-window expects at least one input and init operand"); + return nullptr; + } if (operands.size() % 2) { TokenError(StrCat("expects an even number of operands, but has ", operands.size(), " operands")); return nullptr; } - if (operands.empty()) { - TokenError("reduce-window expects at least one input and init operand"); - return nullptr; - } if (!maybe_infer_shape([&] { return ShapeInference::InferReduceWindowShape( operands[0]->shape(), operands[1]->shape(), *window, diff --git a/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.cc b/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.cc index a29a10e42d34f0..2e045fdfe8f7f2 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.cc +++ b/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.cc @@ -186,12 +186,11 @@ std::vector AsyncCollectiveCreator::MatchCollectives( const HloOpcode op = instruction->opcode(); // We only care about collective ops and collective fusions here. - if (HloPredicateIsNotOp( - instruction)) { + if (HloPredicateIsNotOp< + HloOpcode::kAllReduce, HloOpcode::kAllGather, HloOpcode::kAllToAll, + HloOpcode::kCollectiveBroadcast, HloOpcode::kCollectivePermute, + HloOpcode::kCollectiveReduce, HloOpcode::kReduceScatter, + HloOpcode::kRaggedAllToAll, HloOpcode::kFusion>(instruction)) { continue; } @@ -219,6 +218,10 @@ std::vector AsyncCollectiveCreator::MatchCollectives( bool convert = config_.convert_collective_broadcast(instruction); VLOG(2) << "kCollectiveBroadcast: convert=" << convert; matched = convert; + } else if (op == HloOpcode::kCollectiveReduce) { + bool convert = config_.convert_collective_reduce(instruction); + VLOG(2) << "kCollectiveReduce: convert=" << convert; + matched = convert; } else if (op == HloOpcode::kCollectivePermute) { bool convert = config_.convert_collective_permute(instruction); VLOG(2) << "kCollectivePermute: convert=" << convert; diff --git a/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.h b/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.h index dc5e5f0ec6e24b..a82ad51e16ae35 100644 --- a/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.h +++ b/third_party/xla/xla/hlo/transforms/collectives/async_collective_creator.h @@ -50,6 +50,7 @@ class AsyncCollectiveCreator : public HloModulePass { HloPredicate convert_reduce_scatter = HloPredicateFalse; HloPredicate convert_ragged_all_to_all = HloPredicateFalse; HloPredicate convert_collective_fusion = HloPredicateFalse; + HloPredicate convert_collective_reduce = HloPredicateFalse; HloPredicate should_ignore_size_check = HloPredicateFalse; ContextShapeQuery get_context_shapes = [](const HloInstruction*) { return std::vector{}; diff --git a/third_party/xla/xla/python/ifrt_proxy/client/executable.cc b/third_party/xla/xla/python/ifrt_proxy/client/executable.cc index 46a55a35938b88..b16c106054156b 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/executable.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/executable.cc @@ -805,7 +805,10 @@ LoadedExecutable::Execute(absl::Span args, if (result_needs_exec_status) { req->set_result_status_handle(status_handle); } - rpc_helper_->LoadedExecutableExecute(std::move(req)); + + // Fire and forget the RPC. + (void)rpc_helper_->LoadedExecutableExecute(std::move(req)); + if (result_needs_exec_status) { // Note that the RPCs within `FetchExecuteResult` need to be sent after // `LoadedExecutableExecute` above, or the server will not recognize the diff --git a/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.cc b/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.cc index e89d7fcda6828e..3ca001e1c1610c 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.cc +++ b/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.cc @@ -452,6 +452,17 @@ tsl::Future> DoRpc(RpcHelper::Batcher* batcher, &IfrtResponse::has_##PROPERTY##_response, std::move(req), #PROPERTY); \ } +// RPCs that do not need to return a response. +// TODO(madthanu): For such RPCs, remove the proto response message and stop +// sending a message from the server. +#define RPC_FIRE_FORGET(METHOD, PROPERTY) \ + void RpcHelper::METHOD(std::unique_ptr req) { \ + (void)DoRpc( \ + batcher_.get(), &IfrtRequest::set_allocated_##PROPERTY##_request, \ + &IfrtResponse::mutable_##PROPERTY##_response, \ + &IfrtResponse::has_##PROPERTY##_response, std::move(req), #PROPERTY); \ + } + RPC(Init, init); RPC(GetDefaultDeviceAssignment, get_default_device_assignment); RPC(CheckFuture, check_future); @@ -462,14 +473,15 @@ RPC(MakeErrorArrays, make_error_arrays); RPC(AssembleArrayFromSingleDeviceArrays, assemble_array_from_single_device_arrays); RPC(RemapArrays, remap_arrays); -RPC(ReshardArrays, reshard_arrays); +RPC_FIRE_FORGET(ReshardArrays, reshard_arrays); RPC(BitcastArrays, bitcast_arrays); -RPC(DisassembleIntoSingleDeviceArrays, disassemble_into_single_device_arrays); +RPC_FIRE_FORGET(DisassembleIntoSingleDeviceArrays, + disassemble_into_single_device_arrays); RPC(CopyToHostBuffer, copy_to_host_buffer); RPC(IsArrayDeleted, is_array_deleted); RPC(DestructArray, destruct_array) -RPC(CopyArrays, copy_arrays); -RPC(FullyReplicatedShard, fully_replicated_shard); +RPC_FIRE_FORGET(CopyArrays, copy_arrays); +RPC_FIRE_FORGET(FullyReplicatedShard, fully_replicated_shard); RPC(DeleteArray, delete_array); RPC(Compile, compile); RPC(LoadedExecutableMetadata, loaded_executable_metadata); diff --git a/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.h b/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.h index 247fba01fb626a..bff0bdeb4018b8 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.h +++ b/third_party/xla/xla/python/ifrt_proxy/client/rpc_helper.h @@ -122,17 +122,13 @@ class RpcHelper { std::unique_ptr req); ResponseFuture BitcastArrays( std::unique_ptr req); - ResponseFuture ReshardArrays( - std::unique_ptr req); - ResponseFuture - DisassembleIntoSingleDeviceArrays( + void ReshardArrays(std::unique_ptr req); + void DisassembleIntoSingleDeviceArrays( std::unique_ptr req); ResponseFuture CopyToHostBuffer( std::unique_ptr req); - ResponseFuture CopyArrays( - std::unique_ptr req); - ResponseFuture FullyReplicatedShard( - std::unique_ptr req); + void CopyArrays(std::unique_ptr req); + void FullyReplicatedShard(std::unique_ptr req); ResponseFuture IsArrayDeleted( std::unique_ptr req); ResponseFuture DeleteArray( diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc index b8997931708ff2..01b7531f11c307 100644 --- a/third_party/xla/xla/service/buffer_assignment.cc +++ b/third_party/xla/xla/service/buffer_assignment.cc @@ -531,6 +531,7 @@ absl::Status GatherComputationsByAllocationType( case HloOpcode::kCustomCall: case HloOpcode::kAllReduce: case HloOpcode::kReduceScatter: + case HloOpcode::kCollectiveReduce: case HloOpcode::kAllReduceStart: case HloOpcode::kMap: case HloOpcode::kReduce: diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index dbf5d073be1a37..f0bb4c89cb1d3d 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -606,11 +606,8 @@ cc_library( srcs = ["ir_emitter2.cc"], hdrs = ["ir_emitter2.h"], deps = [ - ":backend_config_proto_cc", ":dot_op_emitter", - ":elemental_ir_emitter", ":ir_emitter", - ":parallel_loop_emitter", "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", @@ -624,10 +621,8 @@ cc_library( "//xla/service:buffer_assignment", "//xla/service:hlo_module_config", "//xla/service/llvm_ir:dynamic_update_slice_util", - "//xla/service/llvm_ir:fused_ir_emitter", "//xla/service/llvm_ir:ir_array", "//xla/service/llvm_ir:llvm_util", - "//xla/service/llvm_ir:loop_emitter", "//xla/stream_executor:launch_dim", "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", @@ -787,23 +782,6 @@ cc_library( ], ) -cc_library( - name = "parallel_loop_emitter", - srcs = ["parallel_loop_emitter.cc"], - hdrs = ["parallel_loop_emitter.h"], - deps = [ - ":ir_emission_utils", - "//xla/service/llvm_ir:ir_array", - "//xla/service/llvm_ir:llvm_loop", - "//xla/service/llvm_ir:llvm_util", - "//xla/service/llvm_ir:loop_emitter", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/strings:str_format", - "@com_google_absl//absl/strings:string_view", - "@llvm-project//llvm:Core", - ], -) - cc_library( name = "thunk_emitter", srcs = ["thunk_emitter.cc"], @@ -835,7 +813,6 @@ cc_library( "//xla/backends/cpu/codegen:target_machine_features", "//xla/backends/cpu/codegen/dot:dot_kernel_emitter", "//xla/backends/cpu/codegen/elemental:concatenate_kernel_emitter", - "//xla/backends/cpu/codegen/elemental:elemental_kernel_emitter", "//xla/backends/cpu/codegen/emitters:cpu_fusion_emitters", "//xla/backends/cpu/runtime:all_gather_thunk", "//xla/backends/cpu/runtime:all_reduce_thunk", @@ -1057,7 +1034,6 @@ xla_cc_test( tags = ["not_run:arm"], deps = [ ":cpu_instruction_fusion", - ":cpu_options", "//xla:literal_util", "//xla:shape_util", "//xla:xla_data_proto_cc", @@ -1083,14 +1059,12 @@ cc_library( srcs = ["cpu_instruction_fusion.cc"], hdrs = ["cpu_instruction_fusion.h"], deps = [ - ":cpu_options", "//xla:shape_util", "//xla:xla_data_proto_cc", "//xla/codegen/emitters:elemental_hlo_to_mlir", "//xla/hlo/analysis:hlo_reachability", "//xla/hlo/ir:hlo", "//xla/service:fusion_node_indexing_evaluation", - "//xla/service:hlo_module_config", "//xla/service:instruction_fusion", "//xla/service:pattern_matcher", "@com_google_absl//absl/algorithm:container", @@ -1148,7 +1122,6 @@ cc_library( "//xla/backends/cpu/codegen:target_machine_features", "//xla/hlo/ir:hlo", "@com_google_absl//absl/log:check", - "@llvm-project//llvm:Core", ], ) diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index b58ee9ea86ddf4..ae23f202029bdd 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -1089,11 +1089,7 @@ absl::Status CpuCompiler::RunHloPassesAfterLayoutAssn( &alias_info, /*may_duplicate=*/!use_multi_output_fusion); - bool use_experimental_loop_fusion = - options::UseExperimentalLoopFusion(module->config()); - bool use_tiled_emitter = options::EnableTiledEmitter(module->config()); - pipeline.AddPass(use_experimental_loop_fusion, - use_tiled_emitter, target_machine_features); + pipeline.AddPass(target_machine_features); if (use_multi_output_fusion) { pipeline.AddPass(&alias_info); @@ -1139,8 +1135,7 @@ absl::Status CpuCompiler::RunHloPassesAfterLayoutAssn( // Safeguard for late elemental instructions created during post-layout // simplification. - pipeline.AddPass(use_experimental_loop_fusion, - use_tiled_emitter, target_machine_features); + pipeline.AddPass(target_machine_features); // Outline ops in the entry computation into calls to subcomputations. if (!is_aot_compile) { @@ -1814,6 +1809,13 @@ CpuCompiler::CompileCpuExecutable( llvm_module.get(), std::move(ir_compiler)); } + TargetMachineFeatures target_machine_features(target_machine.get()); + + // ThunkEmitter needs elemental ops in fusions. Always run the idempotent + // wrapper: run_hlo_passes=false or --xla_disable_hlo_passes may skip it. + FusionWrapper fusion_wrapper(&target_machine_features); + ABSL_RETURN_IF_ERROR(fusion_wrapper.Run(module.get()).status()); + absl::flat_hash_map instruction_to_profile_idx; absl::flat_hash_map @@ -1862,8 +1864,6 @@ CpuCompiler::CompileCpuExecutable( return cpu_executable; }; - TargetMachineFeatures target_machine_features(target_machine.get()); - // TODO(ezhulenev): Once we fully migrate to Thunks current IrEmitter should // be renamed to NestedIrEmitter and be used only for emitting nested (aka // thread local or embedded) computations (reductions, maps, etc.). diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc index 9374b7565201c7..9cdad81fee7e18 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc @@ -65,12 +65,9 @@ TEST_F(CpuCompilerInternalsTest, DotRootedLoopFusionRoutesToMlirEmitter) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, ParseAndReturnVerifiedModule(kDotFusionHlo)); - if (!options::UseExperimentalLoopFusion(hlo_module->config())) { - GTEST_SKIP() << "Test covers the new-fusion-emitter configuration."; - } const auto* fusion = Cast( hlo_module->entry_computation()->root_instruction()); - EXPECT_TRUE(FusionRoutesToMlirEmitter(hlo_module->config(), fusion)); + EXPECT_TRUE(FusionRoutesToMlirEmitter(fusion)); } TEST_F(CpuCompilerInternalsTest, LoopFusionRoutesToMlirEmitter) { @@ -87,12 +84,9 @@ TEST_F(CpuCompilerInternalsTest, LoopFusionRoutesToMlirEmitter) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, ParseAndReturnVerifiedModule(kLoopFusionHlo)); - if (!options::UseExperimentalLoopFusion(hlo_module->config())) { - GTEST_SKIP() << "Test covers the new-fusion-emitter configuration."; - } const auto* fusion = Cast( hlo_module->entry_computation()->root_instruction()); - EXPECT_TRUE(FusionRoutesToMlirEmitter(hlo_module->config(), fusion)); + EXPECT_TRUE(FusionRoutesToMlirEmitter(fusion)); } std::optional GetMetadataInt(llvm::Metadata* absl_nullable value) { diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_test.cc index 00ae3fb4a280c7..5a04fc8f9dfd8d 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler_test.cc @@ -165,6 +165,22 @@ TEST_F(CpuCompilerTest, PermutationSortConvertedToScatter) { EXPECT_TRUE(Run(hlo, /*run_hlo_passes=*/true)); } +TEST_F(CpuCompilerTest, CompilesElementalOpsWithFusionWrapperPassDisabled) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(R"( + HloModule test + ENTRY main { + p0 = f32[8]{0} parameter(0) + p1 = f32[8]{0} parameter(1) + ROOT add = f32[8]{0} add(p0, p1) + } + )")); + module->mutable_config().mutable_debug_options().add_xla_disable_hlo_passes( + "fusion-wrapper"); + + EXPECT_TRUE(Run(std::move(module), /*run_hlo_passes=*/true)); +} + } // namespace } // namespace cpu } // namespace xla diff --git a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc index c34b0af2b71fd8..233312b5887d3f 100644 --- a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc +++ b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc @@ -28,9 +28,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/layout_util.h" -#include "xla/service/cpu/cpu_options.h" #include "xla/service/fusion_node_indexing_evaluation.h" -#include "xla/service/hlo_module_config.h" #include "xla/service/instruction_fusion.h" #include "xla/service/pattern_matcher.h" #include "xla/shape_util.h" @@ -148,24 +146,16 @@ bool IsCoupledReductionShiftExpProducer(const HloInstruction* instr) { } // Should we block the fusion of the subcomputation of the passed instruction? -bool BlockSubcomputationFusion(const HloInstruction* instruction, - const HloModuleConfig& config) { +bool BlockSubcomputationFusion(const HloInstruction* instruction) { HloOpcode opcode = instruction->opcode(); if (opcode == HloOpcode::kScatter) { return true; } - const bool use_experimental_fusion_emitters = - options::UseExperimentalLoopFusion(config); // If the instruction itself can be fused then the subcomputation should be // blocked as the fusion emitter can't emit fusion ops inside another // fusion. - if (use_experimental_fusion_emitters && - emitters::IsSupportedElementalOp(opcode)) { - return true; - } - - return false; + return emitters::IsSupportedElementalOp(opcode); } } // namespace @@ -386,7 +376,7 @@ void CpuInstructionFusion::ComputeInstructionsToSkip( for (HloInstruction* instr : callable->called_computation()->instructions()) instructions_to_skip_.insert(instr); - } else if (BlockSubcomputationFusion(instruction, module->config())) { + } else if (BlockSubcomputationFusion(instruction)) { for (const auto* computation : instruction->called_computations()) { for (const auto* instr : computation->instructions()) { instructions_to_skip_.insert(instr); diff --git a/third_party/xla/xla/service/cpu/cpu_instruction_fusion_test.cc b/third_party/xla/xla/service/cpu/cpu_instruction_fusion_test.cc index b8e4737866a003..29e1e22cd295b5 100644 --- a/third_party/xla/xla/service/cpu/cpu_instruction_fusion_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_instruction_fusion_test.cc @@ -35,7 +35,6 @@ limitations under the License. #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/utils/hlo_matchers.h" #include "xla/literal_util.h" -#include "xla/service/cpu/cpu_options.h" #include "xla/service/transpose_folding.h" #include "xla/shape.h" #include "xla/tests/test_utils.h" @@ -592,10 +591,6 @@ TEST_F(OpcodeFusionTest, DynamicSliceWithDynamicUpdateSlice) { TEST_F(OpcodeFusionTest, MessOfFusibleNodes) { auto module = CreateNewVerifiedModule(); - if (options::UseExperimentalLoopFusion(module->config())) { - GTEST_SKIP() << "New fusion emitter does not support DUS yet."; - } - HloComputation::Builder builder(TestName()); Shape full_shape = ShapeUtil::MakeShape(F32, {4, 100, 10, 100, 50}); diff --git a/third_party/xla/xla/service/cpu/cpu_options.cc b/third_party/xla/xla/service/cpu/cpu_options.cc index d2750b7577a73b..1ae6afe2561722 100644 --- a/third_party/xla/xla/service/cpu/cpu_options.cc +++ b/third_party/xla/xla/service/cpu/cpu_options.cc @@ -155,12 +155,6 @@ std::optional> LlvmIrGemmTileSize( tile_size_n_in_vector_width); } -bool UseExperimentalLoopFusion(const HloModuleConfig& config) { - const auto& extra_options_map = - config.debug_options().xla_backend_extra_options(); - return extra_options_map.count(kDisableNewFusionEmitters) == 0; -} - bool FlattenAfterFusion(const HloModuleConfig& config) { const auto& extra_options_map = config.debug_options().xla_backend_extra_options(); @@ -168,9 +162,6 @@ bool FlattenAfterFusion(const HloModuleConfig& config) { } bool UseMultiOutputFusion(const HloModuleConfig& config) { - if (!options::UseExperimentalLoopFusion(config)) { - return false; - } const auto& extra_options_map = config.debug_options().xla_backend_extra_options(); return extra_options_map.count(kUseMultiOutputFusion) > 0; diff --git a/third_party/xla/xla/service/cpu/cpu_options.h b/third_party/xla/xla/service/cpu/cpu_options.h index de761e9d5b2f8a..911530d22cb06e 100644 --- a/third_party/xla/xla/service/cpu/cpu_options.h +++ b/third_party/xla/xla/service/cpu/cpu_options.h @@ -44,8 +44,6 @@ inline constexpr absl::string_view kFoldAllConstants = "xla_cpu_fold_all_constants"; inline constexpr absl::string_view kSmallWhileLoopByteThreshold = "xla_cpu_small_while_loop_byte_threshold"; -inline constexpr absl::string_view kDisableNewFusionEmitters = - "xla_cpu_disable_new_fusion_emitters"; inline constexpr absl::string_view kFlattenAfterFusion = "xla_cpu_flatten_after_fusion"; inline constexpr absl::string_view kUseMultiOutputFusion = @@ -71,7 +69,6 @@ std::optional> LlvmIrGemmTileSize( const HloModuleConfig& config); absl::StatusOr SmallWhileLoopByteThreshold( const HloModuleConfig& config); -bool UseExperimentalLoopFusion(const HloModuleConfig& config); bool FlattenAfterFusion(const HloModuleConfig& config); bool UseMultiOutputFusion(const HloModuleConfig& config); bool EnableTiledEmitter(const HloModuleConfig& config); diff --git a/third_party/xla/xla/service/cpu/fusion_wrapper.cc b/third_party/xla/xla/service/cpu/fusion_wrapper.cc index a2756acfa2fe7b..dee487ce288224 100644 --- a/third_party/xla/xla/service/cpu/fusion_wrapper.cc +++ b/third_party/xla/xla/service/cpu/fusion_wrapper.cc @@ -99,7 +99,7 @@ bool FusionWrapper::MustWrapInstruction(const HloInstruction& instruction) { case HloOpcode::kTan: case HloOpcode::kTanh: case HloOpcode::kXor: - return using_new_fusion_emitter_; + return true; case HloOpcode::kCopy: // If it is a simple copy with no change in layout then it is more // efficient to use the default copy thunk which will just be a simple diff --git a/third_party/xla/xla/service/cpu/fusion_wrapper.h b/third_party/xla/xla/service/cpu/fusion_wrapper.h index f0d3689f8cd10a..77bd79a4b44f05 100644 --- a/third_party/xla/xla/service/cpu/fusion_wrapper.h +++ b/third_party/xla/xla/service/cpu/fusion_wrapper.h @@ -30,10 +30,8 @@ namespace cpu { class FusionWrapper : public emitters::FusionWrapperBase { public: explicit FusionWrapper( - bool using_new_fusion_emitter, bool use_tiled_emitter, const TargetMachineFeatures* target_machine_features = nullptr) - : using_new_fusion_emitter_(using_new_fusion_emitter), - target_machine_features_(target_machine_features) {} + : target_machine_features_(target_machine_features) {} ~FusionWrapper() override = default; absl::string_view name() const override { return "fusion-wrapper"; } @@ -41,7 +39,6 @@ class FusionWrapper : public emitters::FusionWrapperBase { bool MustWrapInstruction(const HloInstruction& instruction) override; private: - bool using_new_fusion_emitter_; const TargetMachineFeatures* target_machine_features_; }; diff --git a/third_party/xla/xla/service/cpu/fusion_wrapper_test.cc b/third_party/xla/xla/service/cpu/fusion_wrapper_test.cc index c4447f089c1196..80561724855eb5 100644 --- a/third_party/xla/xla/service/cpu/fusion_wrapper_test.cc +++ b/third_party/xla/xla/service/cpu/fusion_wrapper_test.cc @@ -65,8 +65,7 @@ TEST_F(FusionWrapperTest, Scatter) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); EXPECT_TRUE(changed); @@ -76,8 +75,7 @@ TEST_F(FusionWrapperTest, Scatter) { } TEST_F(FusionWrapperTest, TransposeWrappedWithNewFusionEmitters) { - // Standalone transposes route to ElementalKernelEmitter when unwrapped. - // Wrap them when the new fusion emitters are enabled. + // Standalone transposes are wrapped into loop fusions. static constexpr absl::string_view hlo_string = R"( HloModule m ENTRY e { @@ -87,8 +85,7 @@ TEST_F(FusionWrapperTest, TransposeWrappedWithNewFusionEmitters) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); EXPECT_TRUE(changed); EXPECT_EQ(m->entry_computation()->root_instruction()->opcode(), @@ -111,8 +108,7 @@ TEST_F(FusionWrapperTest, DynamicUpdateSliceWrappedWithNewFusionEmitters) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); EXPECT_TRUE(changed); EXPECT_EQ(m->entry_computation()->root_instruction()->opcode(), @@ -133,8 +129,7 @@ TEST_F(FusionWrapperTest, MissingElementalOpcodesWrappedWithNewFusionEmitters) { op); ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/false); + FusionWrapper wrapper; ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); EXPECT_TRUE(changed) << "Failed for opcode: " << op; EXPECT_EQ(m->entry_computation()->root_instruction()->opcode(), @@ -153,8 +148,7 @@ TEST_F(FusionWrapperTest, MissingElementalOpcodesWrappedWithNewFusionEmitters) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/false); + FusionWrapper wrapper; ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); EXPECT_TRUE(changed) << "Failed for opcode: mulhi"; EXPECT_EQ(m->entry_computation()->root_instruction()->opcode(), @@ -162,45 +156,6 @@ TEST_F(FusionWrapperTest, MissingElementalOpcodesWrappedWithNewFusionEmitters) { } } -TEST_F(FusionWrapperTest, - MissingElementalOpcodesNotWrappedWithoutNewFusionEmitters) { - static constexpr absl::string_view kUnaryOpcodes[] = { - "acos", "acosh", "asin", "asinh", "atanh", "cosh", "sinh"}; - for (absl::string_view op : kUnaryOpcodes) { - std::string hlo_string = absl::StrFormat(R"( - HloModule m - ENTRY e { - p0 = f32[64,32] parameter(0) - ROOT r = f32[64,32] %s(p0) - } - )", - op); - ASSERT_OK_AND_ASSIGN(std::unique_ptr m, - ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false); - ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); - EXPECT_FALSE(changed) << "Failed for opcode: " << op; - } - - { - static constexpr absl::string_view hlo_string = R"( - HloModule m - ENTRY e { - p0 = s32[64,32] parameter(0) - p1 = s32[64,32] parameter(1) - ROOT r = s32[64,32] mulhi(p0, p1) - } - )"; - ASSERT_OK_AND_ASSIGN(std::unique_ptr m, - ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false); - ASSERT_OK_AND_ASSIGN(bool changed, wrapper.Run(m.get())); - EXPECT_FALSE(changed) << "Failed for opcode: mulhi"; - } -} - TEST_F(FusionWrapperTest, CopyWithMismatchedLayoutsWrappedWithNewFusionEmitters) { static constexpr absl::string_view hlo_string = R"( @@ -212,8 +167,7 @@ TEST_F(FusionWrapperTest, )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/true, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_TRUE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -228,8 +182,7 @@ TEST_F(FusionWrapperTest, CopyWithMatchingLayoutsNotWrapped) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/true, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_FALSE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -245,8 +198,7 @@ TEST_F(FusionWrapperTest, )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/true, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_TRUE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -262,8 +214,7 @@ TEST_F(FusionWrapperTest, ConcatenateWithMismatchedLayoutsWrapped) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_TRUE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -279,8 +230,7 @@ TEST_F(FusionWrapperTest, ConcatenateWithMatchingLayoutsNotWrapped) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_FALSE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -298,8 +248,7 @@ TEST_F(FusionWrapperTest, NonEigenConvolutionWrappedWithNewFusionEmitters) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_TRUE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -317,8 +266,7 @@ TEST_F(FusionWrapperTest, EigenConvolutionNotWrapped) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_FALSE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } @@ -336,8 +284,7 @@ TEST_F(FusionWrapperTest, NonEigenConvolutionWrapped) { )"; ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(hlo_string)); - FusionWrapper wrapper(/*using_new_fusion_emitter=*/false, - /*use_tiled_emitter=*/false, &target_machine_features_); + FusionWrapper wrapper(&target_machine_features_); EXPECT_TRUE( wrapper.MustWrapInstruction(*m->entry_computation()->root_instruction())); } diff --git a/third_party/xla/xla/service/cpu/ir_emission_utils.h b/third_party/xla/xla/service/cpu/ir_emission_utils.h index b1523d74dda58b..ce1b12c44f41dd 100644 --- a/third_party/xla/xla/service/cpu/ir_emission_utils.h +++ b/third_party/xla/xla/service/cpu/ir_emission_utils.h @@ -17,10 +17,7 @@ limitations under the License. #define XLA_SERVICE_CPU_IR_EMISSION_UTILS_H_ #include -#include -#include -#include "llvm/IR/Value.h" #include "xla/backends/cpu/codegen/target_machine_features.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -40,18 +37,6 @@ bool CanUseEigenConvolution( int64_t GetMinimumAlignmentForArray( const Shape& shape, const TargetMachineFeatures& target_machine_features); -// Dynamic loop bounds are specified as an array of dimension index -// [start, limit) pairs of ir values (one for each partitioned outer dimension). -// -// EX: Let 'shape' = [8, 16, 32], with the loop bounds of the two-most major -// dimensions dynamic. Then 'dynamic_loop_bounds' will contain the -// following ir values for the two most-major dimensions: -// [dim0_index_start_ir_value, dim0_index_limit_ir_value] -// [dim1_index_start_ir_value, dim1_index_limit_ir_value] -// -// See IrFunction and ParallelLoopEmitter for details. -using DynamicLoopBounds = std::vector>; - } // namespace cpu } // namespace xla diff --git a/third_party/xla/xla/service/cpu/ir_emitter2.cc b/third_party/xla/xla/service/cpu/ir_emitter2.cc index af1f3382085e33..61ced9599866d3 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter2.cc +++ b/third_party/xla/xla/service/cpu/ir_emitter2.cc @@ -17,9 +17,7 @@ limitations under the License. #include #include -#include #include -#include #include #include #include @@ -35,14 +33,8 @@ limitations under the License. #include "absl/strings/str_join.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" -#include "absl/types/span.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/Twine.h" #include "llvm/IR/Attributes.h" -#include "llvm/IR/Constants.h" -#include "llvm/IR/DerivedTypes.h" -#include "llvm/IR/GlobalValue.h" -#include "llvm/IR/GlobalVariable.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/MDBuilder.h" @@ -64,19 +56,13 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" -#include "xla/service/cpu/backend_config.pb.h" #include "xla/service/cpu/dot_op_emitter.h" -#include "xla/service/cpu/elemental_ir_emitter.h" #include "xla/service/cpu/ir_emitter.h" -#include "xla/service/cpu/parallel_loop_emitter.h" #include "xla/service/hlo_module_config.h" #include "xla/service/llvm_ir/dynamic_update_slice_util.h" -#include "xla/service/llvm_ir/fused_ir_emitter.h" #include "xla/service/llvm_ir/ir_array.h" #include "xla/service/llvm_ir/llvm_util.h" -#include "xla/service/llvm_ir/loop_emitter.h" #include "xla/shape.h" -#include "xla/shape_partition.h" #include "xla/stream_executor/launch_dim.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" @@ -129,10 +115,6 @@ IrEmitter2::IrEmitter2(const HloModule& hlo_module, llvm::Module* module, KernelApiIrBuilder::Options::FromHloModuleConfig( hlo_module_.config())) {} -bool IrEmitter2::fast_min_max() const { - return hlo_module_.config().debug_options().xla_cpu_enable_fast_min_max(); -} - bool IrEmitter2::IsSupportedByFusionEmitter( const HloFusionInstruction* fusion) const { FusionEmitterKind fusion_emitter_kind = AnalyzeHloFusion(fusion); @@ -190,57 +172,6 @@ absl::StatusOr IrEmitter2::EmitPadHostKernel( KernelInfo(std::move(kernel_prototype), se::BlockDim(), se::ThreadDim())); } -absl::StatusOr IrEmitter2::EmitFusionHostKernel( - const HloFusionInstruction* fusion) { - VLOG(2) << "Emit fusion host kernel: " << fusion->name(); - - // In XLA:CPU output fusion can only be a fusion into dot operation. - if (fusion->fusion_kind() == HloInstruction::FusionKind::kOutput) { - return EmitDotFusionHostKernel(fusion); - } - - if (fusion->fusion_kind() != HloInstruction::FusionKind::kLoop) { - return Internal("Unsupported fusion kind for instruction: %s", - fusion->ToString()); - } - - ABSL_ASSIGN_OR_RETURN(KernelPrototype kernel_prototype, - EmitKernelPrototype(fusion)); - - llvm::IRBuilder<> b(module_->getContext()); - b.SetInsertPoint(kernel_prototype.function->getEntryBlock().getTerminator()); - - IrEmitter::IRBuilderGuard builder_guard = nested_ir_emitter_->WithBuilder(b); - - HloComputation* nested_computation = fusion->fused_instructions_computation(); - ABSL_RETURN_IF_ERROR(nested_ir_emitter_ - ->EmitNestedComputation(*nested_computation, - llvm_ir::IrName(fusion), false) - .status()); - - CpuElementalIrEmitter elemental_emitter = ElementalIrEmmiterFactory(&b); - - FusedIrEmitter fused_emitter(elemental_emitter); - for (int i = 0; i < fusion->operand_count(); i++) { - fused_emitter.BindGenerator( - *fusion->fused_parameter(i), [&, i](llvm_ir::IrArray::Index idx) { - return kernel_prototype.arguments[i].EmitReadArrayElement(idx, &b); - }); - } - - // Emit plain elemental loops for the fusion operation. - ABSL_ASSIGN_OR_RETURN( - auto element_generator, - fused_emitter.GetGenerator(*fusion->fused_expression_root())); - - ABSL_ASSIGN_OR_RETURN( - se::ThreadDim thread_dims, - EmitElementalLoops(b, fusion, kernel_prototype, element_generator)); - - return kernels_.emplace_back( - KernelInfo(std::move(kernel_prototype), se::BlockDim(), thread_dims)); -} - // Dot (fusion) host kernel only supports strategies that emit LLVM IR. static bool IsDotCodegenStrategy(DotImplementationStrategy strategy) { static std::array kDotCodegenStrategies = { @@ -380,139 +311,6 @@ absl::StatusOr IrEmitter2::EmitKernelPrototype( *module_, instr, &nested_ir_emitter_->assignment(), "ir_emitter2"); } -std::optional IrEmitter2::GetParallelConfig( - const HloInstruction* instr) { - // Check if the instruction is marked for parallel execution. - auto backend_config = instr->backend_config(); - if (!backend_config.ok() || - backend_config->outer_dimension_partitions().empty()) { - return std::nullopt; - } - - ParallelConfig config; - config.outer_dimension_partitions.assign( - backend_config->outer_dimension_partitions().begin(), - backend_config->outer_dimension_partitions().end()); - - return config; -} - -IrEmitter2::ParallelPartitionBounds IrEmitter2::EmitParallelPartitionBounds( - llvm::IRBuilderBase& b, const KernelPrototype& kernel_prototype, - const ParallelConfig& parallel_config, const Shape& shape, - absl::string_view name) { - ShapePartitionIterator it(shape, parallel_config.outer_dimension_partitions); - - size_t num_parallel_dimensions = - parallel_config.outer_dimension_partitions.size(); - - // Create a constant array of all partition bounds. We will be indexing into - // this array using block and thread dimension indices passed in a call frame. - // - // Type: [#partitions x [#outer_dimensions x [lower_bound, upper_bound]]] - // - llvm::ArrayType* dim_bounds_ty = llvm::ArrayType::get(b.getInt64Ty(), 2); - llvm::ArrayType* partition_bounds_ty = - llvm::ArrayType::get(dim_bounds_ty, num_parallel_dimensions); - llvm::ArrayType* parallel_bounds_ty = - llvm::ArrayType::get(partition_bounds_ty, it.GetTotalPartitionCount()); - - // Build a nested array of partition bounds from shape partition iterator. - std::vector partition_bounds; - for (int64_t i = 0; i < it.GetTotalPartitionCount(); ++i) { - std::vector dim_counts; - for (auto [lower, size] : it.GetPartition(i)) { - dim_counts.push_back(llvm::ConstantArray::get( - dim_bounds_ty, {b.getInt64(lower), b.getInt64(lower + size)})); - } - partition_bounds.push_back( - llvm::ConstantArray::get(partition_bounds_ty, dim_counts)); - } - - llvm::Constant* parallel_bounds = - llvm::ConstantArray::get(parallel_bounds_ty, partition_bounds); - - llvm::Module* module = b.GetInsertBlock()->getParent()->getParent(); - llvm::GlobalVariable* parallel_bounds_global = new llvm::GlobalVariable( - /*M=*/*module, - /*Ty=*/parallel_bounds_ty, - /*isConstant=*/true, - /*Linkage=*/llvm::GlobalValue::PrivateLinkage, - /*Initializer=*/parallel_bounds, - /*Name=*/absl::StrCat(name, "_parallel_bounds")); - - // Construct IR to load bounds for all parallel dimensions. - ParallelPartitionBounds bounds; - for (size_t i = 0; i < num_parallel_dimensions; ++i) { - llvm::Value* partition = kernel_prototype.workgroup_id.x; - llvm::Value* parallel_dim = b.getInt32(i); - - llvm::Value* lower_gep = b.CreateInBoundsGEP( - parallel_bounds_ty, parallel_bounds_global, - {b.getInt32(0), partition, parallel_dim, b.getInt32(0)}, - absl::StrCat("lo_dim_", i, "_gep")); - - llvm::Value* upper_gep = b.CreateInBoundsGEP( - parallel_bounds_ty, parallel_bounds_global, - {b.getInt32(0), partition, parallel_dim, b.getInt32(1)}, - absl::StrCat("up_dim_", i, "_gep")); - - bounds.emplace_back( - b.CreateLoad(b.getInt64Ty(), lower_gep, absl::StrCat("lo_dim_", i)), - b.CreateLoad(b.getInt64Ty(), upper_gep, absl::StrCat("up_dim_", i))); - } - - return bounds; -} - -absl::StatusOr IrEmitter2::EmitElementalLoops( - llvm::IRBuilderBase& b, const HloInstruction* instr, - const KernelPrototype& kernel_prototype, - const llvm_ir::ElementGenerator& element_generator) { - // We can emit loops for instruction with multiple results only if it is a - // fusion, reduce or reduce window. - bool multiple_results = kernel_prototype.results.size() > 1; - bool support_multiple_results = instr->opcode() == HloOpcode::kFusion || - instr->opcode() == HloOpcode::kReduce || - instr->opcode() == HloOpcode::kReduceWindow; - - auto parallel_config = GetParallelConfig(instr); - bool has_parallel_config = parallel_config.has_value(); - - if (multiple_results && !support_multiple_results) { - return Internal( - "Multi-output host kernels are not supported for %s instruction", - HloOpcodeString(instr->opcode())); - } - - // TODO(ezhulenev): Support multiple results for parallel loops. - if (multiple_results) { - ABSL_RETURN_IF_ERROR( - llvm_ir::LoopEmitter(element_generator, kernel_prototype.results, &b) - .EmitLoop(llvm_ir::IrName(instr))); - return se::ThreadDim(); - } - - const llvm_ir::IrArray& result = kernel_prototype.results.front(); - - // Emit a loop for a single parallel partition with dynamic bounds computed - // from thread index. - if (has_parallel_config) { - ParallelPartitionBounds parallel_bounds = EmitParallelPartitionBounds( - b, kernel_prototype, *parallel_config, instr->shape(), instr->name()); - ABSL_RETURN_IF_ERROR( - ParallelLoopEmitter(element_generator, result, ¶llel_bounds, &b) - .EmitLoop(llvm_ir::IrName(instr))); - return se::ThreadDim(ShapePartitionAssigner::GetTotalPartitionCount( - parallel_config->outer_dimension_partitions)); - } - - // Emit a whole loop for the instruction. - ABSL_RETURN_IF_ERROR(llvm_ir::LoopEmitter(element_generator, result, &b) - .EmitLoop(llvm_ir::IrName(instr))); - return se::ThreadDim(); -} - // This is a convenience function taken from IrEmitter, it uses module_ class // field. If there will be more functions that use module_, we should consider // refactoring (like we did for compute_function_ and builder_). @@ -526,18 +324,4 @@ void IrEmitter2::AttachInvariantLoadMetadataForLoad( hlo_module_.config()); } -CpuElementalIrEmitter IrEmitter2::ElementalIrEmmiterFactory( - llvm::IRBuilderBase* b) const { - auto thread_local_call_fn = [this](const HloComputation& callee, - absl::Span parameters, - absl::string_view name, bool is_reducer) { - return nested_ir_emitter_->EmitThreadLocalCall( - callee, parameters, name, is_reducer, - /*in_compute_function=*/false); - }; - - return CpuElementalIrEmitter(module_, b, thread_local_call_fn, true, - fast_min_max()); -} - } // namespace xla::cpu diff --git a/third_party/xla/xla/service/cpu/ir_emitter2.h b/third_party/xla/xla/service/cpu/ir_emitter2.h index 090c2af66b0c5c..ffbdb1b00b444d 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter2.h +++ b/third_party/xla/xla/service/cpu/ir_emitter2.h @@ -17,9 +17,7 @@ limitations under the License. #define XLA_SERVICE_CPU_IR_EMITTER2_H_ #include -#include #include -#include #include #include "absl/container/flat_hash_set.h" @@ -28,18 +26,14 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "llvm/IR/BasicBlock.h" -#include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/Module.h" -#include "llvm/IR/Value.h" #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/service/buffer_assignment.h" -#include "xla/service/cpu/elemental_ir_emitter.h" #include "xla/service/cpu/ir_emitter.h" -#include "xla/service/llvm_ir/loop_emitter.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/launch_dim.h" @@ -113,10 +107,6 @@ class IrEmitter2 { // Emits a host kernel for the pad instruction. absl::StatusOr EmitPadHostKernel(const HloInstruction* pad); - // Emits a host kernel for the given fusion instruction. - absl::StatusOr EmitFusionHostKernel( - const HloFusionInstruction* fusion); - // Emits a host kernel for the given dot fusion instruction (output fusion). absl::StatusOr EmitDotFusionHostKernel( const HloFusionInstruction* fusion); @@ -137,39 +127,6 @@ class IrEmitter2 { absl::StatusOr EmitKernelPrototype( const HloInstruction* instr); - // Parallel partition bounds for parallelized outer dimensions: - // vector<[i64 lower_bound, i64 upper_bound]> - using ParallelPartitionBounds = - std::vector>; - - // A config for running kernel in parallel. We rely on partitioning iteration - // space along the outer dimension(s) and run each partition as a separate - // task inside a runtime-managed thread pool. - struct ParallelConfig { - std::vector outer_dimension_partitions; - }; - - // Returns parallel config for the given instruction or std::nullopt if - // the instruction has to be compiled to a single threaded loop. - std::optional GetParallelConfig(const HloInstruction* instr); - - // Emits LLVM IR that computes parallel partition bounds from the call frame's - // block and thread dimensions and parallel execution config. - ParallelPartitionBounds EmitParallelPartitionBounds( - llvm::IRBuilderBase& b, const KernelPrototype& kernel_prototype, - const ParallelConfig& parallel_config, const Shape& shape, - absl::string_view name); - - // Emits LLVM IR using elemental loop emitter and the given element generator. - // If the instruction is parallelized, it will emit a parallel loop partition - // and return the requested number of execution threads. - absl::StatusOr EmitElementalLoops( - llvm::IRBuilderBase& b, const HloInstruction* instr, - const KernelPrototype& kernel_prototype, - const llvm_ir::ElementGenerator& element_generator); - - bool fast_min_max() const; - // Returns the number of bytes within the shape. int64_t ByteSizeOf(const Shape& shape) const; @@ -177,8 +134,6 @@ class IrEmitter2 { // load metadata. void AttachInvariantLoadMetadataForLoad(llvm::LoadInst* instr) const; - CpuElementalIrEmitter ElementalIrEmmiterFactory(llvm::IRBuilderBase* b) const; - const HloModule& hlo_module_; llvm::Module* module_; diff --git a/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc b/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc deleted file mode 100644 index 2bfffd88df937e..00000000000000 --- a/third_party/xla/xla/service/cpu/parallel_loop_emitter.cc +++ /dev/null @@ -1,95 +0,0 @@ -/* Copyright 2017 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/service/cpu/parallel_loop_emitter.h" - -#include -#include -#include - -#include "absl/log/check.h" -#include "absl/strings/str_format.h" -#include "absl/strings/string_view.h" -#include "xla/service/llvm_ir/llvm_loop.h" -#include "xla/service/llvm_ir/llvm_util.h" - -namespace xla { -namespace cpu { - -ParallelLoopEmitter::ParallelLoopEmitter( - const llvm_ir::ElementGenerator& target_element_generator, - const llvm_ir::IrArray& target_array, - const DynamicLoopBounds* dynamic_loop_bounds, llvm::IRBuilderBase* b) - : LoopEmitter(target_element_generator, target_array, b), - dynamic_loop_bounds_(dynamic_loop_bounds) {} - -std::vector -ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock(absl::string_view loop_name, - llvm::Type* index_type, - llvm::Value* base_index) { - CHECK_NE(index_type, nullptr); - - CHECK_EQ(base_index, nullptr) - << "XLA CPU implementation of" - << " ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock doesn't support" - << " base_index, but it was requested."; - - CHECK(!shape_.IsTuple()); - CHECK(!ShapeUtil::IsScalar(shape_)); - - llvm_ir::ForLoopNest loop_nest(loop_name, b_); - const int64_t num_dims = shape_.dimensions().size(); - std::vector array_multi_index(num_dims); - - // Add loops from outer-most to inner-most dimensions. - for (int i = LayoutUtil::MinorToMajor(shape_).size() - 1; i >= 0; --i) { - const int64_t dimension = LayoutUtil::Minor(shape_.layout(), i); - const int bounds_index = num_dims - 1 - i; - // Only unroll the most minor dimension, this seems to give us good runtime - // performance with a large improvement in compile time. - auto unroll_mode = (i == 0) ? llvm_ir::UnrollMode::kDefaultUnroll - : llvm_ir::UnrollMode::kNoUnroll; - if (bounds_index < dynamic_loop_bounds_->size()) { - // Emit dynamic loop bounds for this dimension. Dynamic loop bounds - // are read from ir function dynamic loop bounds argument. - llvm::Value* start_index = (*dynamic_loop_bounds_)[bounds_index].first; - llvm::Value* end_index = (*dynamic_loop_bounds_)[bounds_index].second; - - std::unique_ptr loop = loop_nest.AddLoop( - /*suffix=*/absl::StrFormat("dim.%d", dimension), start_index, - end_index, unroll_mode); - array_multi_index[dimension] = loop->GetIndVarValue(); - } else { - // Emit static loop bounds for this dimension. - std::unique_ptr loop = loop_nest.AddLoop( - /*start_index=*/0, - /*end_index=*/shape_.dimensions(dimension), - /*suffix=*/absl::StrFormat("dim.%d", dimension), unroll_mode); - array_multi_index[dimension] = loop->GetIndVarValue(); - } - } - // Point IR builder at inner loop BB. - llvm_ir::SetToFirstInsertPoint(loop_nest.GetInnerLoopBodyBasicBlock(), b_); - - // Set exit_bb_ to the exit block of the loop nest. - exit_bb_ = loop_nest.GetOuterLoopExitBasicBlock(); - CHECK(exit_bb_ != nullptr); - - llvm_ir::IrArray::Index array_index(array_multi_index, shape_, index_type); - return {array_index}; -} - -} // namespace cpu -} // namespace xla diff --git a/third_party/xla/xla/service/cpu/parallel_loop_emitter.h b/third_party/xla/xla/service/cpu/parallel_loop_emitter.h deleted file mode 100644 index 924544f1abcd42..00000000000000 --- a/third_party/xla/xla/service/cpu/parallel_loop_emitter.h +++ /dev/null @@ -1,77 +0,0 @@ -/* Copyright 2017 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ -#define XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ - -#include - -#include "absl/strings/string_view.h" -#include "llvm/IR/IRBuilder.h" -#include "llvm/IR/Value.h" -#include "xla/service/cpu/ir_emission_utils.h" -#include "xla/service/llvm_ir/ir_array.h" -#include "xla/service/llvm_ir/loop_emitter.h" - -namespace xla { -namespace cpu { - -// ParallelLoopEmitter emits a loop nest for the target array shape. -// The outer loop bounds of the loop nest are passed as ir values at runtime -// (specified in 'dynamic_loop_bounds'), and the inner loop bounds are static. -// Dynamic loop bounds are specified as an array of dimension index -// [start, limit) pairs of ir values (one for each partitioned outer dimension). -// -// EX: Let 'shape' = [8, 16, 32], with the loop bounds of the two-most major -// dimensions dynamic. Then 'dynamic_loop_bounds' will contain the -// following ir values for the two most-major dimensions: -// [dim0_index_start_ir_value, dim0_index_limit_ir_value] -// [dim1_index_start_ir_value, dim1_index_limit_ir_value] -// -// Code emitted by ParallelLoopEmitter will be called in a multi-threaded -// context where each thread will be assigned a different set of outer dimension -// partitions, and where all threads will collectively iterate over the -// entire target array shape. -// -// Outer dimension partitions can be generated using the ShapePartitionAssigner -// and ShapePartitionIterator utility classes from shape_partition.cc. -// -class ParallelLoopEmitter : public llvm_ir::LoopEmitter { - public: - // Constructs a ParallelLoopEmitter which uses 'target_element_generator' to - // generate elements, 'dynamic_loop_bounds' to set the loop bounds of the - // most-major dimensions, and 'target_array.' shape to set the static loop - // bounds for the most-minor dimensions. - ParallelLoopEmitter(const llvm_ir::ElementGenerator& target_element_generator, - const llvm_ir::IrArray& target_array, - const DynamicLoopBounds* dynamic_loop_bounds, - llvm::IRBuilderBase* b); - - ParallelLoopEmitter(const ParallelLoopEmitter&) = delete; - ParallelLoopEmitter& operator=(const ParallelLoopEmitter&) = delete; - ~ParallelLoopEmitter() override = default; - - std::vector EmitIndexAndSetExitBasicBlock( - absl::string_view loop_name, llvm::Type* index_type, - llvm::Value* base_index) override; - - private: - const DynamicLoopBounds* dynamic_loop_bounds_; -}; - -} // namespace cpu -} // namespace xla - -#endif // XLA_SERVICE_CPU_PARALLEL_LOOP_EMITTER_H_ diff --git a/third_party/xla/xla/service/cpu/tests/BUILD b/third_party/xla/xla/service/cpu/tests/BUILD index 631d56b3679fe8..3c4427008c24b6 100644 --- a/third_party/xla/xla/service/cpu/tests/BUILD +++ b/third_party/xla/xla/service/cpu/tests/BUILD @@ -104,22 +104,6 @@ xla_cc_test( ], ) -xla_cc_test( - name = "cpu_external_constants_test", - srcs = ["cpu_external_constants_test.cc"], - deps = [ - ":cpu_pjrt_codegen_test", - "//xla:array2d", - "//xla:literal_util", - "//xla:shape_util", - "//xla:xla_data_proto_cc", - "//xla/hlo/ir:hlo", - "//xla/service/cpu:cpu_options", - "//xla/tests:xla_internal_test_main", - "//xla/tsl/platform:test", - ], -) - xla_test( name = "cpu_noalias_test", srcs = ["cpu_noalias_test.cc"], @@ -395,8 +379,6 @@ xla_test( backends = ["cpu"], deps = [ "//xla:literal", - "//xla/service/cpu:fusion_wrapper", - "//xla/service/cpu:target_machine_features_stub", "//xla/tests:hlo_test_base", "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", diff --git a/third_party/xla/xla/service/cpu/tests/cpu_copy_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_copy_test.cc index cb2a62dbc0aabb..f0340e73c37216 100644 --- a/third_party/xla/xla/service/cpu/tests/cpu_copy_test.cc +++ b/third_party/xla/xla/service/cpu/tests/cpu_copy_test.cc @@ -20,8 +20,6 @@ limitations under the License. #include #include "absl/types/span.h" #include "xla/literal.h" -#include "xla/service/cpu/fusion_wrapper.h" -#include "xla/service/cpu/target_machine_features_stub.h" #include "xla/tests/hlo_test_base.h" #include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" @@ -41,12 +39,6 @@ ENTRY entry { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TargetMachineFeaturesStub target_machine_features( - [](int64_t size) { return 16; }); - FusionWrapper fusion_wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/true, - &target_machine_features); - ASSERT_OK(fusion_wrapper.Run(module.get())); ASSERT_OK_AND_ASSIGN(const Literal result, Execute(std::move(module), {}, /*run_hlo_passes=*/false)); @@ -71,12 +63,6 @@ ENTRY entry { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); - TargetMachineFeaturesStub target_machine_features( - [](int64_t size) { return 16; }); - FusionWrapper fusion_wrapper(/*using_new_fusion_emitter=*/true, - /*use_tiled_emitter=*/true, - &target_machine_features); - ASSERT_OK(fusion_wrapper.Run(module.get())); ASSERT_OK_AND_ASSIGN(const Literal result, Execute(std::move(module), {}, /*run_hlo_passes=*/false)); diff --git a/third_party/xla/xla/service/cpu/tests/cpu_external_constants_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_external_constants_test.cc deleted file mode 100644 index b54aed6f98f39a..00000000000000 --- a/third_party/xla/xla/service/cpu/tests/cpu_external_constants_test.cc +++ /dev/null @@ -1,74 +0,0 @@ -/* Copyright 2018 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include -#include - -#include "xla/array2d.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/literal_util.h" -#include "xla/service/cpu/cpu_options.h" -#include "xla/service/cpu/tests/cpu_pjrt_codegen_test.h" -#include "xla/shape_util.h" -#include "xla/tsl/platform/test.h" -#include "xla/xla_data.pb.h" - -namespace xla::cpu { -namespace { - -class CpuExternalConstantsTest : public CpuPjRtCodegenTest { - public: - void TestWithArray(int64_t rows, int64_t cols, - const char* filecheck_pattern) { - HloComputation::Builder builder(TestName()); - - Array2D backing_array(rows, cols); - backing_array.FillUnique(); - - auto shape = ShapeUtil::MakeShape(F32, {rows, cols}); - - HloInstruction* constant = - builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2FromArray2D(backing_array))); - HloInstruction* param = - builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); - builder.AddInstruction( - HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, constant)); - - std::unique_ptr module = CreateNewVerifiedModule(); - if (options::UseExperimentalLoopFusion(module->config())) { - return; - } - - module->AddEntryComputation(builder.Build()); - - CompileAndVerifyIr(std::move(module), filecheck_pattern, - /*match_optimized_ir=*/false); - } -}; - -TEST_F(CpuExternalConstantsTest, DoNotExternalizeConstants) { - TestWithArray(/*rows=*/4, /*cols=*/4, R"( -CHECK-NOT: external unnamed_addr constant [16 x float] -CHECK: @[[CST:.+]] = private unnamed_addr constant [64 x i8] {{.*}}, align {{[0-9]+}} -)"); -} - -} // namespace -} // namespace xla::cpu diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.cc b/third_party/xla/xla/service/cpu/thunk_emitter.cc index 23937c970fc2e9..ba2b102225063f 100644 --- a/third_party/xla/xla/service/cpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc @@ -42,7 +42,6 @@ limitations under the License. #include "xla/backends/cpu/codegen/computation_kernel_emitter.h" #include "xla/backends/cpu/codegen/dot/dot_kernel_emitter.h" #include "xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.h" -#include "xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h" #include "xla/backends/cpu/codegen/fusion_compiler.h" #include "xla/backends/cpu/codegen/fusion_emitter.h" #include "xla/backends/cpu/codegen/ir_compiler.h" @@ -319,73 +318,6 @@ absl::StatusOr ThunkEmitter::EmitHloInstruction( case HloOpcode::kBatchNormTraining: return EmitBatchNormTrainingThunk(instruction); - // Simple HLO instructions lowered to elemental host kernels (plain loops - // behind the HostKernel API). - case HloOpcode::kAbs: - case HloOpcode::kAcos: - case HloOpcode::kAcosh: - case HloOpcode::kAsin: - case HloOpcode::kAsinh: - case HloOpcode::kAdd: - case HloOpcode::kAnd: - case HloOpcode::kAtan2: - case HloOpcode::kAtanh: - case HloOpcode::kBroadcast: - case HloOpcode::kBitcastConvert: - case HloOpcode::kCbrt: - case HloOpcode::kCeil: - case HloOpcode::kClamp: - case HloOpcode::kClz: - case HloOpcode::kCompare: - case HloOpcode::kComplex: - case HloOpcode::kConvert: - case HloOpcode::kCos: - case HloOpcode::kCosh: - case HloOpcode::kDivide: - case HloOpcode::kDynamicUpdateSlice: - case HloOpcode::kErf: - case HloOpcode::kExp: - case HloOpcode::kExpm1: - case HloOpcode::kFloor: - case HloOpcode::kGather: - case HloOpcode::kImag: - case HloOpcode::kIota: - case HloOpcode::kIsFinite: - case HloOpcode::kLog1p: - case HloOpcode::kLog: - case HloOpcode::kMap: - case HloOpcode::kMaximum: - case HloOpcode::kMinimum: - case HloOpcode::kMultiply: - case HloOpcode::kMulhi: - case HloOpcode::kNegate: - case HloOpcode::kNot: - case HloOpcode::kOr: - case HloOpcode::kPopulationCount: - case HloOpcode::kPower: - case HloOpcode::kReal: - case HloOpcode::kReducePrecision: - case HloOpcode::kRemainder: - case HloOpcode::kReshape: - case HloOpcode::kReverse: - case HloOpcode::kRoundNearestAfz: - case HloOpcode::kRoundNearestEven: - case HloOpcode::kRsqrt: - case HloOpcode::kSelect: - case HloOpcode::kShiftLeft: - case HloOpcode::kShiftRightArithmetic: - case HloOpcode::kShiftRightLogical: - case HloOpcode::kSign: - case HloOpcode::kSin: - case HloOpcode::kSinh: - case HloOpcode::kSqrt: - case HloOpcode::kSubtract: - case HloOpcode::kTan: - case HloOpcode::kTanh: - case HloOpcode::kTranspose: - case HloOpcode::kXor: - return EmitElementalKernelThunk(instruction); - // ReplicaId and PartitionId identify the location of the current device in // a logical grid of communicating devices. case HloOpcode::kReplicaId: @@ -407,10 +339,6 @@ absl::StatusOr ThunkEmitter::EmitHloInstruction( case HloOpcode::kPad: return EmitPadKernelThunk(instruction); - case HloOpcode::kSlice: - case HloOpcode::kDynamicSlice: - return EmitSliceThunk(instruction); - case HloOpcode::kConcatenate: return EmitConcatenateKernelThunk(instruction); @@ -439,10 +367,6 @@ absl::StatusOr ThunkEmitter::EmitHloInstruction( } return EmitFusionKernelThunk(instruction); - case HloOpcode::kReduce: - case HloOpcode::kReduceWindow: - return EmitReductionKernelThunk(instruction); - case HloOpcode::kRng: return EmitRngThunk(instruction); @@ -483,9 +407,11 @@ absl::StatusOr ThunkEmitter::EmitHloInstruction( return EmitSortThunk(instruction); default: - return absl::UnimplementedError( - absl::StrCat("HLO opcode `", HloOpcodeString(instruction->opcode()), - "` is not supported by XLA:CPU ThunkEmitter")); + return absl::UnimplementedError(absl::StrCat( + "HLO opcode `", HloOpcodeString(instruction->opcode()), + "` is not supported by XLA:CPU ThunkEmitter. Standalone elemental " + "and reduction operations must be wrapped into fusions by " + "FusionWrapper.")); } } @@ -761,14 +687,7 @@ absl::StatusOr ThunkEmitter::EmitConvolutionThunk( instruction->feature_group_count()); } - // This is a completely un-optimized version of convolution just to - // have an early version that works. E.g. the input index and - // padding calculation is not hoisted out of the inner loop. - // - // See the description of convolution in the XLA documentation for the pseudo - // code for convolution. - VLOG(2) << "Falling back to unoptimized convolution: " << instruction->name(); - return EmitElementalKernelThunk(instruction); + return Unimplemented("Unoptimized convolution is not supported"); } absl::StatusOr ThunkEmitter::EmitCopyThunk( @@ -781,23 +700,6 @@ absl::StatusOr ThunkEmitter::EmitCopyThunk( instruction->shape()); } -absl::StatusOr ThunkEmitter::EmitElementalKernelThunk( - const HloInstruction* instruction) { - ElementalKernelEmitter emitter(instruction, &buffer_assignment_, - &target_machine_features_); - ABSL_ASSIGN_OR_RETURN(KernelDefinition kernel_definition, - emitter.EmitKernelDefinition()); - - auto kernel_spec = kernel_definition.spec(); - auto kernel_source = std::move(kernel_definition).TakeSource(); - - kernels_.push_back( - {kernel_spec.name(), std::move(kernel_source).thread_safe_module()}); - - return MakeKernelThunkSequence(instruction, std::move(kernel_spec), - /*min_alignment=*/MinAlign()); -} - absl::StatusOr ThunkEmitter::EmitPadKernelThunk( const HloInstruction* instruction) { const HloPadInstruction* padInstr = Cast(instruction); @@ -812,7 +714,7 @@ absl::StatusOr ThunkEmitter::EmitFusionKernelThunk( const HloInstruction* instruction) { auto* fusion = Cast(instruction); - if (FusionRoutesToMlirEmitter(hlo_module_config_, fusion)) { + if (FusionRoutesToMlirEmitter(fusion)) { ABSL_ASSIGN_OR_RETURN(std::string fingerprint, GetFusionFingerprint(*fusion, buffer_assignment_, GetDefaultBufferAlignment())); @@ -839,26 +741,23 @@ absl::StatusOr ThunkEmitter::EmitFusionKernelThunk( /*min_alignment=*/MinAlign()); } - ABSL_ASSIGN_OR_RETURN(auto kernel, ir_emitter_.EmitFusionHostKernel(fusion)); - ABSL_ASSIGN_OR_RETURN(auto buffers, GetHostKernelAllocationSlices(instruction)); + if (fusion->fusion_kind() == HloInstruction::FusionKind::kOutput) { + ABSL_ASSIGN_OR_RETURN(auto kernel, ir_emitter_.EmitDotFusionHostKernel(fusion)); + ABSL_ASSIGN_OR_RETURN(auto buffers, GetHostKernelAllocationSlices(instruction)); - return MakeKernelThunkSequence(instruction, buffers, kernel, - /*min_alignment=*/MinAlign()); + return MakeKernelThunkSequence(instruction, buffers, kernel, + /*min_alignment=*/MinAlign()); + } + + return Unimplemented("Unsupported fusion instruction: %s", + fusion->ToString()); } -bool FusionRoutesToMlirEmitter(const HloModuleConfig& config, - const HloFusionInstruction* fusion) { +bool FusionRoutesToMlirEmitter(const HloFusionInstruction* fusion) { if (fusion->fused_expression_root()->opcode() == HloOpcode::kScatter) { return true; } - return options::UseExperimentalLoopFusion(config) && - fusion->fusion_kind() == HloFusionInstruction::FusionKind::kLoop; -} - -absl::StatusOr ThunkEmitter::EmitReductionKernelThunk( - const HloInstruction* instruction) { - // TODO(ezhulenev): Port vectorized reduction emitter from IrEmitter. - return EmitElementalKernelThunk(instruction); + return fusion->fusion_kind() == HloFusionInstruction::FusionKind::kLoop; } absl::StatusOr ThunkEmitter::EmitRngThunk( @@ -1267,14 +1166,6 @@ absl::StatusOr ThunkEmitter::EmitSliceToDynamicThunk( /*min_alignment=*/MinAlign()); } -absl::StatusOr ThunkEmitter::EmitSliceThunk( - const HloInstruction* instruction) { - // TODO(ezhulenev): Consider implementing slice operations as separate - // Thunks because it might be easier to get peak performance from hand - // written code (Eigen slice expression for example). - return EmitElementalKernelThunk(instruction); -} - // Parse the sort comparator to determine the sort direction. std::optional ThunkEmitter::MatchSortDirection( const HloComputation* hlo_comparator) const { diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.h b/third_party/xla/xla/service/cpu/thunk_emitter.h index 7f837199310cc7..1d6f4693657c16 100644 --- a/third_party/xla/xla/service/cpu/thunk_emitter.h +++ b/third_party/xla/xla/service/cpu/thunk_emitter.h @@ -133,9 +133,6 @@ class ThunkEmitter { absl::StatusOr EmitCopyThunk( const HloInstruction* instruction); - absl::StatusOr EmitElementalKernelThunk( - const HloInstruction* instruction); - absl::StatusOr EmitPadKernelThunk( const HloInstruction* instruction); @@ -144,9 +141,6 @@ class ThunkEmitter { absl::StatusOr EmitFusionKernelThunk( const HloInstruction* instruction); - absl::StatusOr EmitReductionKernelThunk( - const HloInstruction* instruction); - absl::StatusOr EmitRngThunk(const HloInstruction* instruction); absl::StatusOr EmitRngBitGeneratorThunk( @@ -208,9 +202,6 @@ class ThunkEmitter { absl::StatusOr EmitOneDnnOpThunk( const HloInstruction* instruction); - absl::StatusOr EmitSliceThunk( - const HloInstruction* instruction); - absl::StatusOr EmitSortThunk( const HloInstruction* instruction); @@ -276,8 +267,7 @@ class ThunkEmitter { ParallelFusionEmitter parallel_fusion_emitter_; }; -bool FusionRoutesToMlirEmitter(const HloModuleConfig& config, - const HloFusionInstruction* fusion); +bool FusionRoutesToMlirEmitter(const HloFusionInstruction* fusion); } // namespace xla::cpu diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index e671287b8aa1ff..6b875d38711fb8 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -514,6 +514,7 @@ cc_library( "//xla/backends/gpu/runtime:collective_broadcast_thunk", "//xla/backends/gpu/runtime:collective_group_thunk", "//xla/backends/gpu/runtime:collective_permute_thunk", + "//xla/backends/gpu/runtime:collective_reduce_thunk", "//xla/backends/gpu/runtime:collective_thunk", "//xla/backends/gpu/runtime:conditional_thunk", "//xla/backends/gpu/runtime:convolution_reorder_thunk", diff --git a/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc b/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc index ab225d333a7762..8ab5a811af3662 100644 --- a/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc +++ b/third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc @@ -897,6 +897,12 @@ absl::Status RunAsyncCollectivesConversionPasses(HloModule* module) { config.convert_all_reduce = HloPredicateTrue; config.convert_all_to_all = HloPredicateTrue; config.convert_collective_broadcast = HloPredicateTrue; + config.convert_collective_reduce = + module->config() + .debug_options() + .xla_gpu_experimental_emit_collective_reduce() + ? HloPredicateTrue + : HloPredicateFalse; config.convert_collective_permute = HloPredicateTrue; config.convert_ragged_all_to_all = HloPredicateTrue; config.convert_reduce_scatter = HloPredicateTrue; @@ -943,6 +949,8 @@ absl::Status RunAsyncCollectivesConversionPasses(HloModule* module) { case HloOpcode::kCollectiveBroadcast: return !disabled_async_ops.contains( DebugOptions::COLLECTIVEBROADCAST); + case HloOpcode::kCollectiveReduce: + return !disabled_async_ops.contains(DebugOptions::ALLREDUCE); case HloOpcode::kReduceScatter: return !disabled_async_ops.contains(DebugOptions::REDUCESCATTER); case HloOpcode::kAllToAll: diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc index 3713b9c37535a1..d6858f394287da 100644 --- a/third_party/xla/xla/service/gpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc @@ -76,6 +76,7 @@ limitations under the License. #include "xla/backends/gpu/runtime/collective_broadcast_thunk.h" #include "xla/backends/gpu/runtime/collective_group_thunk.h" #include "xla/backends/gpu/runtime/collective_permute_thunk.h" +#include "xla/backends/gpu/runtime/collective_reduce_thunk.h" #include "xla/backends/gpu/runtime/collective_thunk.h" #include "xla/backends/gpu/runtime/conditional_thunk.h" #include "xla/backends/gpu/runtime/convolution_reorder_thunk.h" @@ -306,6 +307,7 @@ absl::StatusOr ThunkEmitter::DispatchAsyncDone( case HloOpcode::kRaggedAllToAll: case HloOpcode::kCollectiveBroadcast: case HloOpcode::kCollectivePermute: + case HloOpcode::kCollectiveReduce: return EmitAsyncDone(instr, instr->operand(0)); // Complete a fusion or call wrapped in generic async start/done. @@ -1249,6 +1251,7 @@ class NativeCustomCallEmitterContextImpl absl::StatusOr ThunkEmitter::EmitNativeCustomCallThunks( const HloCustomCallInstruction* instr, NativeCustomCallHandlerRef handler) { + auto released_lock_keeper = llvm_options_lock_->TemporarilyReleaseLock(); NativeCustomCallEmitterContextImpl ctx(this, instr); return handler(*instr, ctx); } @@ -2094,6 +2097,23 @@ Future ThunkEmitter::EmitCollective( Thunk::kCollectiveBroadcast, Cast(collective), std::nullopt); + case HloOpcode::kCollectiveReduce: { + if (!ir_emitter_context_->debug_options() + .xla_gpu_experimental_emit_collective_reduce()) { + return Internal( + "Unsupported collective instruction: %s. Set " + "--xla_gpu_experimental_emit_collective_reduce to enable " + "CollectiveReduce support on GPU.", + collective->ToString()); + } + auto* collective_reduce = + Cast(collective); + return EmitCollective( + Thunk::kCollectiveReduce, collective_reduce, + collective_reduce->use_global_device_ids()); + } + default: return Internal("Unsupported collective instruction: %s", collective->ToString()); @@ -2124,11 +2144,13 @@ Future ThunkEmitter::EmitCollective( << "; partition count: " << partition_count << "; operand count: " << operand_count; - // A collective-broadcast may select its root rank at runtime, in which case - // the last operand is a root-rank vector rather than data to broadcast. + // A collective-broadcast or collective-reduce may select its root rank at run + // time, in which case the last operand is an S32 root-rank vector rather than + // data being broadcast/reduced. const bool has_dynamic_root = [](const HloInstType* inst) { if constexpr (std::is_same_v) { + HloCollectiveBroadcastInstruction> || + std::is_same_v) { return inst->has_dynamic_root(); } return false; @@ -2218,6 +2240,14 @@ Future ThunkEmitter::EmitCollective( info, inst, /*buffers=*/std::move(buffers), ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p(), has_dynamic_root); + } else if constexpr (std::is_same_v) { + // CollectiveReduceThunk needs the dynamic-root flag so it can treat the + // trailing root-rank buffer specially at run time. + thunks = ThunkSequence::Of( + info, inst, /*buffers=*/std::move(buffers), + ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p(), + has_dynamic_root); } else if constexpr (std::is_constructible_v< CollectiveThunkType, Thunk::ThunkInfo, decltype(inst), @@ -2752,6 +2782,7 @@ Future ThunkEmitter::EmitHloInstruction( case HloOpcode::kAllReduce: case HloOpcode::kAllToAll: case HloOpcode::kCollectiveBroadcast: + case HloOpcode::kCollectiveReduce: case HloOpcode::kCollectivePermute: case HloOpcode::kRaggedAllToAll: case HloOpcode::kReduceScatter: diff --git a/third_party/xla/xla/stream_executor/gpu/all_reduce_kernel.h b/third_party/xla/xla/stream_executor/gpu/all_reduce_kernel.h index f116cad0a3cd4e..b029f95cebf4c7 100644 --- a/third_party/xla/xla/stream_executor/gpu/all_reduce_kernel.h +++ b/third_party/xla/xla/stream_executor/gpu/all_reduce_kernel.h @@ -57,7 +57,7 @@ void AbslStringify(Sink& sink, AllReduceStrategy strategy) { // The maximum number of input pointers that can be passed to the all-reduce // kernel. -inline constexpr int64_t kMaxNumAllReduceInputPtrs = 8; +inline constexpr int64_t kMaxNumAllReduceInputPtrs = 64; inline constexpr int64_t kNumElementsPerThread = 4; // A pointer to a buffer that does not alias with other buffers. diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc index b10b9e6a852d5b..ffb756e6b92bce 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.cc @@ -46,8 +46,11 @@ absl::Status IsValidDeviceOrdinal(int device_ordinal, // Returns true if the oneAPI version is 2024.2 or newer. // oneAPI 2024.2 corresponds to __LIBSYCL_MAJOR_VERSION == 7 and // __LIBSYCL_MINOR_VERSION == 2. +// oneAPI 2025.x corresponds to __LIBSYCL_MAJOR_VERSION == 8, and +// oneAPI 2026.x corresponds to __LIBSYCL_MAJOR_VERSION == 9. bool IsOneAPIVersionAtLeast2024_2() { - return (__LIBSYCL_MAJOR_VERSION >= 7) && (__LIBSYCL_MINOR_VERSION >= 2); + return ((__LIBSYCL_MAJOR_VERSION == 7) && (__LIBSYCL_MINOR_VERSION >= 2)) || + (__LIBSYCL_MAJOR_VERSION >= 8); } absl::Status MemcpyDeviceToHost(::sycl::queue* stream_handle, void* dst_host, diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index f85b93f8da690d..670198036147fb 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -3694,6 +3694,7 @@ cc_library( "//xla/service/cpu:executable_proto_cc", "//xla/service/gpu:gpu_executable_proto_cc", "//xla/stream_executor:kernel_spec_proto_cc", + "//xla/stream_executor/abi:executable_abi_version_proto_cc", "//xla/tsl/platform:env", "//xla/tsl/platform:logging", "//xla/util/split_proto:human_readable_aot_executable_proto_cc", diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD index 6b59ed77da4a0f..78530ff65ec847 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD @@ -30,7 +30,6 @@ xla_test( "gb300", ], tags = [ - "broken", "multi_gpu", "no_oss", ], diff --git a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc index 5b1fab26b8c1da..2d2f16502b92aa 100644 --- a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc +++ b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc @@ -45,6 +45,7 @@ limitations under the License. #include "xla/service/device_assignment.h" #include "xla/service/gpu/gpu_executable.pb.h" #include "xla/status_macros.h" +#include "xla/stream_executor/abi/executable_abi_version.pb.h" #include "xla/stream_executor/kernel_spec.pb.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/logging.h" @@ -231,6 +232,15 @@ absl::Status AOTInterceptionPjrtClient::CompareGPUExecutables( ExecutableBuildOptionsProto::descriptor()->FindFieldByName( "debug_options"), HloModuleConfigProto::descriptor()->FindFieldByName("debug_options"), + stream_executor::ExecutableAbiVersionProto::CudaPlatformVersion:: + descriptor() + ->FindFieldByName("cuda_toolkit_version"), + stream_executor::ExecutableAbiVersionProto::CudaPlatformVersion:: + descriptor() + ->FindFieldByName("cudnn_version"), + stream_executor::ExecutableAbiVersionProto::CudaPlatformVersion:: + descriptor() + ->FindFieldByName("cub_version"), }); } diff --git a/third_party/xla/xla/util/split_proto/BUILD b/third_party/xla/xla/util/split_proto/BUILD index 8a5122362547d3..c97b03382cb693 100644 --- a/third_party/xla/xla/util/split_proto/BUILD +++ b/third_party/xla/xla/util/split_proto/BUILD @@ -62,6 +62,7 @@ cc_library( visibility = internal_visibility([ "//learning/brain/research/pjrt:__subpackages__", "//xla:__subpackages__", + "//xla:friends", ]), deps = [ ":split_proto_cc", diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index 44ef782f522cc1..a23060f7f1ee10 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -1032,6 +1032,13 @@ message DebugOptions { // Enable NCCL symmetric buffers. optional bool xla_gpu_experimental_enable_nccl_symmetric_buffers = 406; + // If true, XLA:GPU emits a CollectiveReduceThunk for kCollectiveReduce HLO + // instructions. This is guarded behind a flag to preserve the 2-week + // forward compatibility window between the compiler and the runtime. + // Need to remove this and enable by default once the 2-week + // window has passed. + optional bool xla_gpu_experimental_emit_collective_reduce = 550; + // If true, the GPU latency hiding scheduler selectively overlaps async // device-to-device memcpys with compute-bound kernels. optional bool xla_gpu_experimental_enable_selective_memcpy_overlap = 524; @@ -1875,7 +1882,7 @@ message DebugOptions { // Enables Raft for stable TopK. optional bool xla_gpu_experimental_enable_raft_for_stable_topk = 546; - // Next id: 550 + // Next id: 551 // Extra options to pass to the compilation backend (e.g. LLVM); specific // interpretation of these values is left to the backend.