From 0b69beaf442e3c44434f809410cd680c76aae446 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Sun, 23 Aug 2026 12:06:23 -0400 Subject: [PATCH 01/23] Add missing SoftsignGrad gradient registration Differentiating through tf.nn.softsign twice failed with "LookupError: gradient registry has no entry for: SoftsignGrad", because every other activation's backward op has a Python gradient registration while SoftsignGrad did not. The new registration differentiates the kernel expression gradients / (1 + |features|)^2, giving -2 * gradients * sign(features) / (1 + |features|)^3 for the backprop input and reusing the kernel for the gradients input. Test Plan: Ran softsign_op_test.py against a pip tf-nightly build with the patched nn_grad.py overlaid: "Ran 5 tests in 0.461s / OK (skipped=1)" including the new testGradGrad; with the pristine nn_grad.py that test fails with the LookupError quoted above. Second derivatives verified against -2*sign(x)/(1+|x|)^3 at seven points from -0.9 to 0.9. pylint with tensorflow/tools/ci_build/pylintrc rates both changed Python files 10.00/10. --- RELEASE.md | 6 +++++ .../kernel_tests/nn_ops/softsign_op_test.py | 23 +++++++++++++++++++ tensorflow/python/ops/nn_grad.py | 10 ++++++++ 3 files changed, 39 insertions(+) diff --git a/RELEASE.md b/RELEASE.md index 575ac36ab4b3b6..1c82b90434a66d 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -40,6 +40,12 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr `tf.image.adjust_contrast` can now be differentiated with `GradientTape`. Fixes [#126083](https://github.com/tensorflow/tensorflow/issues/126083). +* `tf.nn.softsign` + + * Fixes second-order gradients of `tf.nn.softsign`. Differentiating twice + previously failed with a lookup error because the `SoftsignGrad` + backward op had no registered Python gradient. + * `tf.experimental.numpy` diff --git a/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py index 585c6435aae524..67f64ffaee2514 100644 --- a/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py @@ -16,6 +16,7 @@ import numpy as np +from tensorflow.python.eager import backprop from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import test_util @@ -25,6 +26,10 @@ from tensorflow.python.platform import test +def _softsign_grad_grad(activation): + return -2 * np.sign(activation) / (1 + np.abs(activation))**3 + + class SoftsignTest(test.TestCase): def _npSoftsign(self, np_features): @@ -82,6 +87,24 @@ def testNoInts(self): "'features' has DataType int32 not in list of allowed values"): nn_ops.softsign(constant_op.constant(7)).eval() + def testGradGrad(self): + with self.cached_session(): + + def f(x): + with backprop.GradientTape(persistent=True) as tape: + tape.watch(x) + y = nn_ops.softsign(x) + dy = tape.gradient(y, x) + return tape.gradient(dy, x) + + x = np.asarray( + [[-0.9, -0.7, -0.5, -0.3, -0.1], [0.1, 0.3, 0.5, 0.7, 0.9]], + dtype=np.float64, + order="F") + got = self.evaluate(f(constant_op.constant(x))) + want = _softsign_grad_grad(x) + self.assertAllClose(got, want) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index a7b4a1c2b63e1f..66a815c73a4087 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -492,6 +492,16 @@ def _SoftsignGrad(op: ops.Operation, grad): return gen_nn_ops.softsign_grad(grad, op.inputs[0]) +@ops.RegisterGradient("SoftsignGrad") +def _SoftsignGradGrad(op: ops.Operation, grad): + x = op.inputs[1] + return ( + gen_nn_ops.softsign_grad(grad, x), + -2.0 * grad * op.inputs[0] * math_ops.sign(x) / + math_ops.pow(1.0 + math_ops.abs(x), 3), + ) + + @ops.RegisterGradient("ReluGrad") def _ReluGradGrad(op: ops.Operation, grad): x = op.inputs[1] From 0cead06e7d4a394e9d8ecd7c94eeeba650327a42 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 12:04:12 -0400 Subject: [PATCH 02/23] SoftsignGrad review cleanup Replace the pow call with three multiplications, which is cheaper and avoids a pow round trip for a small integer exponent, and cover both float32 and float64 in the second derivative test instead of float64 only. Caught in code review on pull request 126092. Test Plan: softsign_op_test.py against the nightly overlay: "OK (skipped=1)". pylint --rcfile=tensorflow/tools/ci_build/pylintrc on both files: 10.00/10. --- .../kernel_tests/nn_ops/softsign_op_test.py | 15 ++++++++------- tensorflow/python/ops/nn_grad.py | 3 ++- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py index 67f64ffaee2514..51c0230d53a66d 100644 --- a/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/softsign_op_test.py @@ -97,13 +97,14 @@ def f(x): dy = tape.gradient(y, x) return tape.gradient(dy, x) - x = np.asarray( - [[-0.9, -0.7, -0.5, -0.3, -0.1], [0.1, 0.3, 0.5, 0.7, 0.9]], - dtype=np.float64, - order="F") - got = self.evaluate(f(constant_op.constant(x))) - want = _softsign_grad_grad(x) - self.assertAllClose(got, want) + for dtype in (np.float32, np.float64): + x = np.asarray( + [[-0.9, -0.7, -0.5, -0.3, -0.1], [0.1, 0.3, 0.5, 0.7, 0.9]], + dtype=dtype, + order="F") + got = self.evaluate(f(constant_op.constant(x))) + want = _softsign_grad_grad(x.astype(np.float64)).astype(dtype) + self.assertAllClose(got, want) if __name__ == "__main__": diff --git a/tensorflow/python/ops/nn_grad.py b/tensorflow/python/ops/nn_grad.py index 66a815c73a4087..5ebe22f0f8fad8 100644 --- a/tensorflow/python/ops/nn_grad.py +++ b/tensorflow/python/ops/nn_grad.py @@ -495,10 +495,11 @@ def _SoftsignGrad(op: ops.Operation, grad): @ops.RegisterGradient("SoftsignGrad") def _SoftsignGradGrad(op: ops.Operation, grad): x = op.inputs[1] + denominator = 1.0 + math_ops.abs(x) return ( gen_nn_ops.softsign_grad(grad, x), -2.0 * grad * op.inputs[0] * math_ops.sign(x) / - math_ops.pow(1.0 + math_ops.abs(x), 3), + (denominator * denominator * denominator), ) From dacb0ffbd02e39d1481f3504940bc2750a1fa91e Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Wed, 26 Aug 2026 03:54:20 -0400 Subject: [PATCH 03/23] Regenerate pywrap_gradient_exclusions for SoftsignGrad The new SoftsignGrad registration changes the gradient exclusion tables that pywrap_gradient_exclusions.cc holds: the registration reads only the grad argument and none of op.inputs, so the op gains a full unused-inputs entry. Regenerated with the documented generator entry point; the output for unmodified master reproduces the committed file byte for byte, and the only delta here is the single new SoftsignGrad line. --- tensorflow/python/eager/pywrap_gradient_exclusions.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index 8c34921427963b..f491195e29214f 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -773,6 +773,7 @@ absl::optional> OpGradientUnusedOutputIndices( {"Softplus"}, {"SoftplusGrad"}, {"Softsign"}, + {"SoftsignGrad"}, {"SpaceToBatch"}, {"SpaceToBatchND"}, {"SpaceToDepth"}, From 32c92fbd62787625d1bcfcf3003693d236b8b62d Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Thu, 27 Aug 2026 09:58:19 -0400 Subject: [PATCH 04/23] Fix OpGradientUnusedOutputIndices array size for SoftsignGrad The regenerated table for the SoftsignGrad registration has 489 entries but the declared std::array size stayed at 488, which fails compilation with excess elements in array initializer. Bump the size to 489. --- tensorflow/python/eager/pywrap_gradient_exclusions.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index f491195e29214f..6e39edcad17aa1 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -429,7 +429,7 @@ absl::optional> OpGradientUnusedInputIndices( absl::optional> OpGradientUnusedOutputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Abs"}, {"AccumulateNV2"}, {"Acos"}, From b67784d0ceed0e465abb1d65259e7c71e2197b5d Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Fri, 28 Aug 2026 01:37:00 -0700 Subject: [PATCH 05/23] Add low-precision support for sort, compare, and select on XLA:CPU Teach CpuFloatSupport and OneDnnFloatSupport that BF16 and F16 are supported for kSort, kCompare, and kSelect. Previously, kSort operands and comparators were unconditionally upcast to F32 during FloatNormalization, preventing low-precision sorting from utilizing optimized inlined sort logic. Also add kSort to FloatSupport::SupportsMixedPrecisions so that key-value sorts with mixed element types (e.g. BF16 keys and S32 values) are preserved. PiperOrigin-RevId: 972429241 --- third_party/xla/xla/service/cpu/BUILD | 4 +- .../xla/xla/service/cpu/cpu_float_support.h | 37 +++++ .../xla/service/cpu/cpu_float_support_test.cc | 127 +++++++++++++++++- .../xla/service/cpu/onednn_float_support.cc | 3 + 4 files changed, 162 insertions(+), 9 deletions(-) diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 59b301065844fb..99dc48e2494892 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -1636,6 +1636,7 @@ cc_library( hdrs = ["cpu_float_support.h"], copts = tsl_copts(), deps = [ + "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/service:float_support", ], @@ -1648,14 +1649,11 @@ xla_cc_test( ":cpu_float_support", "//xla:shape_util", "//xla:xla_data_proto_cc", - "//xla/backends/cpu/codegen:target_machine_features", "//xla/backends/cpu/codegen:target_machine_test_base", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:verified_hlo_module", "//xla/hlo/transforms/simplifiers:float_normalization", "//xla/service:hlo_module_config", - "//xla/tsl/platform:statusor", - "//xla/tsl/platform:test", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/service/cpu/cpu_float_support.h b/third_party/xla/xla/service/cpu/cpu_float_support.h index 2a5ea4670bc31f..e9502839bb008a 100644 --- a/third_party/xla/xla/service/cpu/cpu_float_support.h +++ b/third_party/xla/xla/service/cpu/cpu_float_support.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef XLA_SERVICE_CPU_CPU_FLOAT_SUPPORT_H_ #define XLA_SERVICE_CPU_CPU_FLOAT_SUPPORT_H_ +#include #include #include "xla/hlo/ir/hlo_casting_utils.h" @@ -23,6 +24,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/service/float_support.h" +#include "xla/xla_data.pb.h" namespace xla { namespace cpu { @@ -36,6 +38,41 @@ class CpuFloatSupport : public FloatSupport { : FloatSupport(low_precision_type), call_library_for_instruction_(call_library_for_instruction) {} + bool SupportsLowPrecisionOperand(const HloInstruction& hlo, + int64_t operand_index) const override { + if (LowPrecisionType() == BF16 || LowPrecisionType() == F16) { + switch (hlo.opcode()) { + case HloOpcode::kSort: + case HloOpcode::kCompare: + case HloOpcode::kSelect: + return true; + default: + break; + } + } + return FloatSupport::SupportsLowPrecisionOperand(hlo, operand_index); + } + + bool SupportsLowPrecisionOutput(const HloInstruction& hlo) const override { + if (LowPrecisionType() == BF16 || LowPrecisionType() == F16) { + switch (hlo.opcode()) { + case HloOpcode::kSort: + case HloOpcode::kSelect: + return true; + default: + break; + } + } + return FloatSupport::SupportsLowPrecisionOutput(hlo); + } + + bool SupportsMixedPrecisions(const HloInstruction& hlo) const override { + if (hlo.opcode() == HloOpcode::kSort) { + return true; + } + return FloatSupport::SupportsMixedPrecisions(hlo); + } + // Skip trying to upcast the dot if the dot is supported by a library. bool ShouldSkipInstruction(const HloInstruction& hlo) const override { return (hlo.opcode() == HloOpcode::kDot || diff --git a/third_party/xla/xla/service/cpu/cpu_float_support_test.cc b/third_party/xla/xla/service/cpu/cpu_float_support_test.cc index 9008d97488be48..1891860619ff0f 100644 --- a/third_party/xla/xla/service/cpu/cpu_float_support_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_float_support_test.cc @@ -17,14 +17,13 @@ limitations under the License. #include #include -#include #include -#include "absl/strings/match.h" +#include +#include #include "absl/strings/str_cat.h" #include "absl/strings/str_replace.h" #include "absl/strings/string_view.h" -#include "xla/backends/cpu/codegen/target_machine_features.h" #include "xla/backends/cpu/codegen/target_machine_test_base.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" @@ -33,8 +32,6 @@ limitations under the License. #include "xla/service/hlo_module_config.h" #include "xla/shape.h" #include "xla/shape_util.h" -#include "xla/tsl/platform/statusor.h" -#include "xla/tsl/platform/test.h" #include "xla/xla_data.pb.h" namespace xla::cpu { @@ -153,7 +150,7 @@ TEST_P(SkipInstructionTest, Bf16InF32Out) { // Run FloatNormalization and check the results. FloatNormalization float_normalization(&cpu_float_support); - TF_ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); EXPECT_EQ(upcast, spec.upcast); PrimitiveType expected_input_dtype = spec.upcast ? F32 : BF16; CheckDtype(module.get(), expected_input_dtype, expected_input_dtype, F32); @@ -202,5 +199,123 @@ INSTANTIATE_TEST_SUITE_P(SkipInstructionTestSuite, SkipInstructionTest, ::testing::ValuesIn(GetSkipInstructionTestSpecs()), SkipInstructionTest::Name); +TEST_F(TargetMachineTestBase, SortNotUpcast) { + for (absl::string_view type_str : {"bf16", "f16"}) { + std::string hlo_text = absl::StrReplaceAll(R"( +HloModule test_module + +compare { + p0 = $type$[] parameter(0) + p1 = $type$[] parameter(1) + p2 = s32[] parameter(2) + p3 = s32[] parameter(3) + ROOT cmp = pred[] compare(p0, p1), direction=LT +} + +ENTRY main { + k = $type$[100] parameter(0) + v = s32[100] parameter(1) + ROOT sort = ($type$[100], s32[100]) sort(k, v), dimensions={0}, to_apply=compare, is_stable=true +} +)", + {{"$type$", type_str}}); + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + + PrimitiveType low_precision_type = (type_str == "bf16") ? BF16 : F16; + + CpuFloatSupport cpu_float_support( + low_precision_type, [](const HloInstruction&) { return false; }); + + FloatNormalization float_normalization(&cpu_float_support); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + EXPECT_FALSE(upcast); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kSort); + EXPECT_EQ(root->operand(0)->shape().element_type(), low_precision_type); + EXPECT_EQ(root->operand(1)->shape().element_type(), S32); + + HloComputation* compare_comp = root->to_apply(); + EXPECT_EQ(compare_comp->parameter_instruction(0)->shape().element_type(), + low_precision_type); + } +} + +TEST_F(TargetMachineTestBase, SortUpcastForUnsupportedType) { + std::string hlo_text = R"( +HloModule test_module + +compare { + p0 = f8e5m2[] parameter(0) + p1 = f8e5m2[] parameter(1) + p2 = s32[] parameter(2) + p3 = s32[] parameter(3) + ROOT cmp = pred[] compare(p0, p1), direction=LT +} + +ENTRY main { + k = f8e5m2[100] parameter(0) + v = s32[100] parameter(1) + ROOT sort = (f8e5m2[100], s32[100]) sort(k, v), dimensions={0}, to_apply=compare, is_stable=true +} +)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + + CpuFloatSupport cpu_float_support( + F8E5M2, [](const HloInstruction&) { return false; }); + + FloatNormalization float_normalization(&cpu_float_support); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + EXPECT_TRUE(upcast); + + HloInstruction* sort_instr = FindInstruction(module.get(), HloOpcode::kSort); + ASSERT_NE(sort_instr, nullptr); + EXPECT_EQ(sort_instr->operand(0)->shape().element_type(), F32); + EXPECT_EQ(sort_instr->operand(1)->shape().element_type(), S32); + HloComputation* updated_cmp = sort_instr->to_apply(); + EXPECT_EQ(updated_cmp->parameter_instruction(0)->shape().element_type(), F32); + EXPECT_EQ(updated_cmp->parameter_instruction(1)->shape().element_type(), F32); +} + +TEST_F(TargetMachineTestBase, CompareAndSelectNotUpcast) { + for (absl::string_view type_str : {"bf16", "f16"}) { + std::string hlo_text = absl::StrReplaceAll(R"( +HloModule test_module + +ENTRY main { + p0 = $type$[100] parameter(0) + p1 = $type$[100] parameter(1) + cmp = pred[100] compare(p0, p1), direction=LT + ROOT select = $type$[100] select(cmp, p0, p1) +} +)", + {{"$type$", type_str}}); + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text)); + + PrimitiveType low_precision_type = (type_str == "bf16") ? BF16 : F16; + + CpuFloatSupport cpu_float_support( + low_precision_type, [](const HloInstruction&) { return false; }); + + FloatNormalization float_normalization(&cpu_float_support); + ASSERT_OK_AND_ASSIGN(bool upcast, float_normalization.Run(module.get())); + EXPECT_FALSE(upcast); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kSelect); + EXPECT_EQ(root->shape().element_type(), low_precision_type); + EXPECT_EQ(root->operand(1)->shape().element_type(), low_precision_type); + EXPECT_EQ(root->operand(2)->shape().element_type(), low_precision_type); + + const HloInstruction* cmp = root->operand(0); + EXPECT_EQ(cmp->opcode(), HloOpcode::kCompare); + EXPECT_EQ(cmp->operand(0)->shape().element_type(), low_precision_type); + EXPECT_EQ(cmp->operand(1)->shape().element_type(), low_precision_type); + } +} + } // namespace } // namespace xla::cpu diff --git a/third_party/xla/xla/service/cpu/onednn_float_support.cc b/third_party/xla/xla/service/cpu/onednn_float_support.cc index 84495f0def5b02..b749cccbdd49f6 100644 --- a/third_party/xla/xla/service/cpu/onednn_float_support.cc +++ b/third_party/xla/xla/service/cpu/onednn_float_support.cc @@ -50,6 +50,7 @@ bool OneDnnFloatSupport::IsSupported(const HloInstruction& hlo) const { case HloOpcode::kReshape: case HloOpcode::kReverse: case HloOpcode::kScatter: + case HloOpcode::kCompare: case HloOpcode::kSelect: case HloOpcode::kSelectAndScatter: case HloOpcode::kSlice: @@ -57,6 +58,8 @@ bool OneDnnFloatSupport::IsSupported(const HloInstruction& hlo) const { // Other special ops. case HloOpcode::kBitcast: return true; + case HloOpcode::kSort: + return LowPrecisionType() == BF16 || LowPrecisionType() == F16; default: return false; } From c502330fe4f0e7fb149c62e30cdfa3c832f06cdc Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Fri, 28 Aug 2026 01:37:56 -0700 Subject: [PATCH 06/23] [XLA:GPU] Migrate GpuAotCompilationTest to HloPjRtGpuTestBase Migrate GpuAotCompilationTest from HloTestBaseLegacy to HloPjRtGpuTestBase. Configure AOT compilation options using gpu_target_config() instead of an explicit StreamExecutor, and execute the Triton AOT compiled module via test_runner() PjRt interfaces. PiperOrigin-RevId: 972429647 --- third_party/xla/xla/service/gpu/BUILD | 18 ++- .../service/gpu/gpu_aot_compilation_test.cc | 111 ++++++++---------- 2 files changed, 53 insertions(+), 76 deletions(-) diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index eb3e044fe4fc25..bbe83ba91d3b30 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -2655,16 +2655,14 @@ xla_test( ], ) -xla_cc_test( +xla_test( name = "gpu_aot_compilation_test", srcs = [ "gpu_aot_compilation_test.cc", ], + backends = ["gpu"], tags = [ - "gpu", - "no_oss", "nomsan", # Pulls in precompiled NVIDIA libraries which cause false positives in msan. - "requires-gpu-nvidia", ], deps = if_cuda_is_configured([ ":nvptx_compiler_impl", @@ -2678,26 +2676,24 @@ xla_cc_test( "//xla:literal_util", "//xla:xla_proto_cc", "//xla/backends/gpu/codegen/triton:support", + "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base", "//xla/hlo/ir:hlo", + "//xla/pjrt/proto:compile_options_proto_cc", "//xla/service:compiled_module", "//xla/service:compiler", "//xla/service:executable", - "//xla/service:gpu_plugin", "//xla/service:gpu_topology", "//xla/service:hlo_runner_interface", - "//xla/service:platform_util", - "//xla/stream_executor:platform", - "//xla/stream_executor:platform_manager", - "//xla/stream_executor:stream_executor_h", "//xla/tests:literal_test_util", - "//xla/tests:xla_internal_test_main", # build_cleaner: keep - "//xla/tests/restricted:hlo_test_base_legacy", "//xla/tsl/platform:statusor", + "//xla/util/split_proto:split_executable_and_options_writer", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", + "@riegeli//riegeli/bytes:string_writer", "@tsl//tsl/platform:statusor", ], ) diff --git a/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc b/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc index efff25b0ff66d1..3c6f46c29f779f 100644 --- a/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_aot_compilation_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include #include -#include "absl/strings/ascii.h" #include "absl/strings/escaping.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" @@ -30,22 +29,21 @@ limitations under the License. #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/MLIRContext.h" +#include "riegeli/bytes/string_writer.h" #include "xla/backends/gpu/codegen/triton/support.h" +#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/pjrt/proto/compile_options.pb.h" #include "xla/service/compiled_module.h" #include "xla/service/compiler.h" #include "xla/service/executable.h" #include "xla/service/gpu/gpu_executable.h" #include "xla/service/gpu_topology.h" #include "xla/service/hlo_runner_interface.h" -#include "xla/service/platform_util.h" -#include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_manager.h" -#include "xla/stream_executor/stream_executor.h" #include "xla/tests/literal_test_util.h" -#include "xla/tests/restricted/hlo_test_base_legacy.h" +#include "xla/util/split_proto/split_executable_and_options_writer.h" #include "xla/xla.pb.h" namespace xla { @@ -53,9 +51,12 @@ namespace gpu { using ::testing::IsEmpty; using ::testing::Not; -class GpuAotCompilationTest : public HloTestBaseLegacy { +class GpuAotCompilationTest : public HloPjRtGpuTestBase { protected: - void SetUp() override { debug_options_ = GetDebugOptionsForTest(); } + void SetUp() override { + HloPjRtGpuTestBase::SetUp(); + debug_options_ = GetDebugOptionsForTest(); + } DebugOptions debug_options_; }; @@ -72,35 +73,28 @@ TEST_F(GpuAotCompilationTest, ExportAndLoadExecutable) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); - auto compiler = backend().compiler(); - auto name = - absl::AsciiStrToUpper(PlatformUtil::CanonicalPlatformName("gpu").value()); - ASSERT_OK_AND_ASSIGN(se::Platform * platform, - se::PlatformManager::PlatformWithName(name)); - ASSERT_OK_AND_ASSIGN(se::StreamExecutor * stream_exec, - platform->ExecutorForDevice(0)); - // Compile AOT. - AotCompilationOptions aot_options(compiler->PlatformId()); - aot_options.set_executor(stream_exec); + AotCompilationOptions aot_options(stream_executor_platform_id()); + aot_options.set_gpu_topology( + GetSingleDeviceGpuTopology("", gpu_target_config())); ASSERT_OK_AND_ASSIGN( std::vector> aot_results, - compiler->CompileAheadOfTime(std::move(module), aot_options)); + compiler()->CompileAheadOfTime(std::move(module), aot_options)); // Serialize-deserialize AOT compilation result. ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); ASSERT_OK_AND_ASSIGN( std::unique_ptr aot_result, - compiler->LoadAotCompilationResult(serialized_aot_result)); + compiler()->LoadAotCompilationResult(serialized_aot_result)); // Load Executable from AOT compilation result. ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, std::move(*aot_result) - .LoadExecutable(compiler->PlatformId(), - stream_exec->GetDeviceDescription(), debug_options_)); + .LoadExecutable(stream_executor_platform_id(), device_description(), + debug_options_)); auto* gpu_executable = dynamic_cast(executable.get()); ASSERT_NE(gpu_executable, nullptr); @@ -120,37 +114,28 @@ TEST_F(GpuAotCompilationTest, AotCompilationWithoutGpuDevice) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); - auto compiler = backend().compiler(); - auto name = - absl::AsciiStrToUpper(PlatformUtil::CanonicalPlatformName("gpu").value()); - ASSERT_OK_AND_ASSIGN(se::Platform * platform, - se::PlatformManager::PlatformWithName(name)); - ASSERT_OK_AND_ASSIGN(se::StreamExecutor * stream_exec, - platform->ExecutorForDevice(0)); - // Stream executor is not passed as an option. - Compiler::GpuTargetConfig gpu_target_config(stream_exec); - AotCompilationOptions aot_options(compiler->PlatformId()); + AotCompilationOptions aot_options(stream_executor_platform_id()); aot_options.set_gpu_topology( - GetSingleDeviceGpuTopology("", gpu_target_config)); + GetSingleDeviceGpuTopology("", gpu_target_config())); ASSERT_OK_AND_ASSIGN( std::vector> aot_results, - compiler->CompileAheadOfTime(std::move(module), aot_options)); + compiler()->CompileAheadOfTime(std::move(module), aot_options)); // Serialize-deserialize AOT compilation result. ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); ASSERT_OK_AND_ASSIGN( std::unique_ptr aot_result, - compiler->LoadAotCompilationResult(serialized_aot_result)); + compiler()->LoadAotCompilationResult(serialized_aot_result)); // Load Executable from AOT compilation result. ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, std::move(*aot_result) - .LoadExecutable(compiler->PlatformId(), - stream_exec->GetDeviceDescription(), debug_options_)); + .LoadExecutable(stream_executor_platform_id(), device_description(), + debug_options_)); } namespace { @@ -211,11 +196,8 @@ std::string CreateTritonCustomCallBackendConfig() { } // namespace TEST_F(GpuAotCompilationTest, ExportAndLoadExecutableWithTriton) { - auto triton_support = - EnsureTritonSupportsComputeCapability(backend() - .default_stream_executor() - ->GetDeviceDescription() - .gpu_compute_capability()); + auto triton_support = EnsureTritonSupportsComputeCapability( + device_description().gpu_compute_capability()); if (!triton_support.ok()) { GTEST_SKIP() << triton_support; } @@ -237,45 +219,44 @@ TEST_F(GpuAotCompilationTest, ExportAndLoadExecutableWithTriton) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); - auto compiler = backend().compiler(); - auto platform_name = - absl::AsciiStrToUpper(PlatformUtil::CanonicalPlatformName("gpu").value()); - ASSERT_OK_AND_ASSIGN(se::Platform * platform, - se::PlatformManager::PlatformWithName(platform_name)); - ASSERT_OK_AND_ASSIGN(se::StreamExecutor * stream_exec, - platform->ExecutorForDevice(0)); - // Compile AOT. - AotCompilationOptions aot_options(compiler->PlatformId()); - aot_options.set_executor(stream_exec); + AotCompilationOptions aot_options(stream_executor_platform_id()); + aot_options.set_gpu_topology( + GetSingleDeviceGpuTopology("", gpu_target_config())); ASSERT_OK_AND_ASSIGN( std::vector> aot_results, - compiler->CompileAheadOfTime(std::move(module), aot_options)); + compiler()->CompileAheadOfTime(std::move(module), aot_options)); // Serialize-deserialize AOT compilation result. ASSERT_OK_AND_ASSIGN(std::string serialized_aot_result, aot_results[0]->SerializeAsString()); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr aot_result, - compiler->LoadAotCompilationResult(serialized_aot_result)); - // Load Executable from AOT compilation result. + // Load and execute via PjRt test runner. + ExecutableAndOptionsProto proto; + *proto.mutable_serialized_executable() = std::move(serialized_aot_result); + proto.set_pjrt_client_name("PjRtStreamExecutorClient"); + proto.mutable_compile_options() + ->mutable_executable_build_options() + ->set_num_replicas(1); + proto.mutable_compile_options() + ->mutable_executable_build_options() + ->set_num_partitions(1); + std::string serialized_split_proto; + ASSERT_OK(WriteSplitExecutableAndOptions( + proto, + std::make_unique>(&serialized_split_proto))); ASSERT_OK_AND_ASSIGN( - std::unique_ptr executable, - std::move(*aot_result) - .LoadExecutable(compiler->PlatformId(), - stream_exec->GetDeviceDescription(), debug_options_)); - std::unique_ptr wrapped_executable = - test_runner_as_hlo_runner().WrapExecutable(std::move(executable)); + std::unique_ptr executable, + test_runner().DeserializeExecutable(serialized_split_proto)); const xla::Literal literal_1 = xla::LiteralUtil::CreateR0(1.0f); const xla::Literal literal_2 = xla::LiteralUtil::CreateR0(2.0f); const xla::Literal literal_3 = xla::LiteralUtil::CreateR0(3.0f); ASSERT_OK_AND_ASSIGN(Literal result, - test_runner_as_hlo_runner().ExecuteWithExecutable( - wrapped_executable.get(), {&literal_1, &literal_3})); + test_runner().ExecuteWithExecutable( + executable.get(), {&literal_1, &literal_3})); EXPECT_TRUE(LiteralTestUtil::Equal( LiteralUtil::MakeTuple({&literal_2, &literal_3}), result)); From aacf4d38efa58fffd74e6d750b9790a8003c9929 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:46:02 -0700 Subject: [PATCH 07/23] PR #47908: Bump ml-dtypes from 0.5.4 to 0.6.0 Imported from GitHub PR https://github.com/openxla/xla/pull/47908 Bumps [ml-dtypes](https://github.com/jax-ml/ml_dtypes) from 0.5.4 to 0.6.0.
Release notes

Sourced from ml-dtypes's releases.

v0.6.0

  • Added new 16-bit complex types: ml_dtypes.complex32 (based on float16) and ml_dtypes.bcomplex32 (based on bfloat16) (#351).
  • Added new 1-bit integer types: ml_dtypes.int1 and ml_dtypes.uint1.
  • Added __format__ method to custom float, complex, and integer types (#341). Previously, formatting custom scalars (e.g. in f-strings) fell back to string formatting, which could truncate exponents or fail on numeric format specifiers.
  • ml_dtypes.finfo and ml_dtypes.iinfo now allow passing an array object to their constructor (#350).
  • Fixed equality comparison (==, !=) when comparing custom dtypes against incompatible types like strings or None.
  • Dropped support for Python 3.9, which reached end-of-life in October 2025.
  • Dropped support for Python 3.13 free-threading, because cibuildwheel dropped support.
  • Dropped support for NumPy < 2.0.
  • Added support for Python 3.15.
  • Switched build system to scikit-build-core and CMake (#361).
  • Updated Eigen dependency to v5.0.1.
Changelog

Sourced from ml-dtypes's changelog.

[0.6.0] - 2026-08-13

  • Added new 16-bit complex types: ml_dtypes.complex32 (based on float16) and ml_dtypes.bcomplex32 (based on bfloat16) (#351).
  • Added new 1-bit integer types: ml_dtypes.int1 and ml_dtypes.uint1.
  • Added __format__ method to custom float, complex, and integer types (#341). Previously, formatting custom scalars (e.g. in f-strings) fell back to string formatting, which could truncate exponents or fail on numeric format specifiers.
  • ml_dtypes.finfo and ml_dtypes.iinfo now allow passing an array object to their constructor (#350).
  • Fixed equality comparison (==, !=) when comparing custom dtypes against incompatible types like strings or None.
  • Dropped support for Python 3.9, which reached end-of-life in October 2025.
  • Dropped support for Python 3.13 free-threading, because cibuildwheel dropped support.
  • Dropped support for NumPy < 2.0.
  • Added support for Python 3.15.
  • Switched build system to scikit-build-core and CMake (#361).
  • Updated Eigen dependency to v5.0.1.
Commits
  • 6bc762d Merge pull request #389 from hawkinsp:release
  • ce2c7bf Prepare v0.6.0 release.
  • 66a7982 Merge pull request #390 from hawkinsp:npy2
  • 2c29114 Drop support for NumPy 1.x.
  • 3a8d5da Merge pull request #388 from hawkinsp:build
  • dec45f7 Merge pull request #387 from hawkinsp:format
  • 2e0e13e Clean up scikit-build-core configuration in pyproject.toml.
  • 90bebd5 Implement format for custom float, complex, and int dtypes.
  • f9c6392 Merge pull request #368 from jax-ml:dependabot/github_actions/actions/downloa...
  • 14396b0 Merge pull request #384 from jax-ml:dependabot/github_actions/actions/setup-p...
  • Additional commits viewable in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=ml-dtypes&package-manager=pip&previous-version=0.5.4&new-version=0.6.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Copybara import of the project: -- 2e9317f2988a2dfb8d4cc32f7bd98e8e620e0c88 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>: Bump ml-dtypes from 0.5.4 to 0.6.0 Bumps [ml-dtypes](https://github.com/jax-ml/ml_dtypes) from 0.5.4 to 0.6.0. - [Release notes](https://github.com/jax-ml/ml_dtypes/releases) - [Changelog](https://github.com/jax-ml/ml_dtypes/blob/main/CHANGELOG.md) - [Commits](https://github.com/jax-ml/ml_dtypes/compare/v0.5.4...v0.6.0) --- updated-dependencies: - dependency-name: ml-dtypes dependency-version: 0.6.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Merging this change closes #47908 PiperOrigin-RevId: 972433287 --- third_party/xla/requirements_lock_3_11.txt | 81 +++++++++++----------- third_party/xla/requirements_lock_3_12.txt | 81 +++++++++++----------- 2 files changed, 82 insertions(+), 80 deletions(-) diff --git a/third_party/xla/requirements_lock_3_11.txt b/third_party/xla/requirements_lock_3_11.txt index c130f9c07482a7..761e0ace834bd6 100644 --- a/third_party/xla/requirements_lock_3_11.txt +++ b/third_party/xla/requirements_lock_3_11.txt @@ -73,43 +73,44 @@ numpy==2.4.6 \ --hash=sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20 lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b -ml-dtypes==0.5.4 \ - --hash=sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf \ - --hash=sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d \ - --hash=sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f \ - --hash=sha256:19b9a53598f21e453ea2fbda8aa783c20faff8e1eeb0d7ab899309a0053f1483 \ - --hash=sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7 \ - --hash=sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22 \ - --hash=sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6 \ - --hash=sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175 \ - --hash=sha256:388d399a2152dd79a3f0456a952284a99ee5c93d3e2f8dfe25977511e0515270 \ - --hash=sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1 \ - --hash=sha256:3d277bf3637f2a62176f4575512e9ff9ef51d00e39626d9fe4a161992f355af2 \ - --hash=sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1 \ - --hash=sha256:4ff7f3e7ca2972e7de850e7b8fcbb355304271e2933dd90814c1cb847414d6e2 \ - --hash=sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298 \ - --hash=sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d \ - --hash=sha256:557a31a390b7e9439056644cb80ed0735a6e3e3bb09d67fd5687e4b04238d1de \ - --hash=sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049 \ - --hash=sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d \ - --hash=sha256:6c7ecb74c4bd71db68a6bea1edf8da8c34f3d9fe218f038814fd1d310ac76c90 \ - --hash=sha256:7c23c54a00ae43edf48d44066a7ec31e05fdc2eee0be2b8b50dd1903a1db94bb \ - --hash=sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465 \ - --hash=sha256:88c982aac7cb1cbe8cbb4e7f253072b1df872701fcaf48d84ffbb433b6568f24 \ - --hash=sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453 \ - --hash=sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56 \ - --hash=sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48 \ - --hash=sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff \ - --hash=sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460 \ - --hash=sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac \ - --hash=sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900 \ - --hash=sha256:a9b61c19040397970d18d7737375cffd83b1f36a11dd4ad19f83a016f736c3ef \ - --hash=sha256:b4b801ebe0b477be666696bda493a9be8356f1f0057a57f1e35cd26928823e5a \ - --hash=sha256:b95e97e470fe60ed493fd9ae3911d8da4ebac16bd21f87ffa2b7c588bf22ea2c \ - --hash=sha256:bc11d7e8c44a65115d05e2ab9989d1e045125d7be8e05a071a48bc76eb6d6040 \ - --hash=sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9 \ - --hash=sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7 \ - --hash=sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6 \ - --hash=sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b \ - --hash=sha256:d81fdb088defa30eb37bf390bb7dde35d3a83ec112ac8e33d75ab28cc29dd8b0 \ - --hash=sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328 +ml-dtypes==0.6.0 \ + --hash=sha256:008382aeab529df5d3f00501ad9a7dcd64494d4b5b1971fc4c79019e6c1f5010 \ + --hash=sha256:03ce583adfce34ad33aa9e1fc7a8344dcf90ea776cc4ef0e5a48d4eae84e5d20 \ + --hash=sha256:084dfe51a7ad58b171f05115f8226ed4233a454a1611371947e806e76f0c638d \ + --hash=sha256:26b1f1fa4f0435a2946859823f6e2bf06796f1e9f10f5a05b08a5e3c8f46ff69 \ + --hash=sha256:28d676428b104bb9717b0928bc5c5129f2d6b51b6727587cc4289e7bf8713cb5 \ + --hash=sha256:2a3e9d53925597fbffafd2a37048dadeddd0bdaba58058f6ae0869ed709a184d \ + --hash=sha256:3035518e3e19add1a4cac9236ab22888b208a4074912514313ccb2d6d242cde8 \ + --hash=sha256:317be9967fb84b0ce4e80e6b1bf71213d21971621cf6f1e501a63602a95297bf \ + --hash=sha256:31f1ce979d31a357e95aa81812f20412c8c954fa43c44ee3ead1e1c8a78575ef \ + --hash=sha256:37da32aa97749251025666d62372775019594577b9c9e9cfda83bed48d778fdb \ + --hash=sha256:3b4a480aa8fd54a1805b8ac10f3f91763926a74f73c0c364c10f9231854f4170 \ + --hash=sha256:3be9911d953f97cddded4b9961d7b650473b7e55806d20f6176f8356dfe7b38e \ + --hash=sha256:3e169214e0d80ff1c038e1b3017e33c23e43bdf948d42d31de8283111c7e2fa3 \ + --hash=sha256:488c99ab181a2f59d9ec3b12c5fa11ec904e92be2c4ba18cded54dd7501208fe \ + --hash=sha256:5359c588cc62de6f78d7430f06b65853d884955494d86d6ad90b6dd64a3f3a08 \ + --hash=sha256:573b11f3c327e17ef3826d266e676cf1149a1f3016f822a05f2306c55d8246bf \ + --hash=sha256:57ed0d6b4ac5e7868361303a9c57fbcf63b768236ee14456f585dfcf260d0292 \ + --hash=sha256:5a519c9e95a216fbcb8e759793ef7fb40793fc803ed839142d6dc5be9be5bc89 \ + --hash=sha256:5e60251d32ced5598972e4d5e06a2f044341f9291402551a3f6f0ec44f9299b0 \ + --hash=sha256:6c8e39b53e90afda8ce52859c93de4dba3e02b76d85dcf091cc469f9184c6dae \ + --hash=sha256:6eaed129a4afe90694b8685e2f9b6294849f5eda4af9a15be83a4326eeebd775 \ + --hash=sha256:6ec0d244a5bba12239025389ad88bbfb45f9f10e25ab4f678e9a4768ebd47532 \ + --hash=sha256:7728c0420ec1c338564fc8b01015ff2d58567e70f17fedce5a0a7c0308c0d5b9 \ + --hash=sha256:84fa136b8602c8c39e3b6cb24918960cd6f36cade7a70376f56770729cd56510 \ + --hash=sha256:8f490c003369ce60e514a0c3b12374f05274c101fee1bead6740ec8a564032b0 \ + --hash=sha256:9c6ad60af4102789a5c09824004beade2f7f28cd1cd581ee5c170d9dc2fbb00e \ + --hash=sha256:b1b503864fada3f74fabf8d9fee7b4c1cbe956301e6fdece975d5f77c2fce958 \ + --hash=sha256:b76fa1d3f92967d58289ac47ab7458ede66e6f3527fff3e59142aee57d9307cd \ + --hash=sha256:bad8d1dd5bed060a29332b99d63d0e5c2969081e1c6ea54adfbccfdfa783be44 \ + --hash=sha256:ce7563e0b1a4482cbc1b4a6272145e54e4489e54fe7428f94908c3d87103abfa \ + --hash=sha256:d4f1b9329a251e4affe3bb58f4d3e2db22a714396fd7ffb40d0b5db423c24d17 \ + --hash=sha256:d574c2b28921dc72e869df248f1a278f6eee176a1f237c8642e1a71eb15f3977 \ + --hash=sha256:de9d14748dbf3968951436ef514a29c9d1fe438aa680d110134ee2f7a9f9df18 \ + --hash=sha256:e25bb3b0ad1217b60626e4ed45b10ca170c41d99fbe44a12bebc1e07ec4aad55 \ + --hash=sha256:e2d6149f3a57f405bcad5fb41e03218b8373936253f23e1ca84c0108abbc3392 \ + --hash=sha256:e74266ca8e97874a937b7646378c178025650a236584f7474d10d8086a6edea3 \ + --hash=sha256:f4adb4af61516510d786cf8c01851a66f6d3ddfa79e1144deaa5b40d8507231e \ + --hash=sha256:f4f59f83c82ab480e924b988e7b1b4eb4de836dfcf5390c6f59148d1a00e1d02 \ + --hash=sha256:f6cb525101b6b903779188c1e9e9490c343b455ab822883e02cf01e5547338d2 \ + --hash=sha256:fb87f46b4f7ad7b5d3ad8f4b452b024bd4229d44c8ff934798c1fe656210387a diff --git a/third_party/xla/requirements_lock_3_12.txt b/third_party/xla/requirements_lock_3_12.txt index e6fb1e433c71f7..097af0257f1d12 100644 --- a/third_party/xla/requirements_lock_3_12.txt +++ b/third_party/xla/requirements_lock_3_12.txt @@ -45,43 +45,44 @@ numpy==2.5.1 \ --hash=sha256:f7feb014281029e628ba2d5a007407443b06e418b6fe451d1e2adcbc8eba0107 lit==17.0.6 \ --hash=sha256:dfa9af9b55fc4509a56be7bf2346f079d7f4a242d583b9f2e0b078fd0abae31b -ml-dtypes==0.5.4 \ - --hash=sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf \ - --hash=sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d \ - --hash=sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f \ - --hash=sha256:19b9a53598f21e453ea2fbda8aa783c20faff8e1eeb0d7ab899309a0053f1483 \ - --hash=sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7 \ - --hash=sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22 \ - --hash=sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6 \ - --hash=sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175 \ - --hash=sha256:388d399a2152dd79a3f0456a952284a99ee5c93d3e2f8dfe25977511e0515270 \ - --hash=sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1 \ - --hash=sha256:3d277bf3637f2a62176f4575512e9ff9ef51d00e39626d9fe4a161992f355af2 \ - --hash=sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1 \ - --hash=sha256:4ff7f3e7ca2972e7de850e7b8fcbb355304271e2933dd90814c1cb847414d6e2 \ - --hash=sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298 \ - --hash=sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d \ - --hash=sha256:557a31a390b7e9439056644cb80ed0735a6e3e3bb09d67fd5687e4b04238d1de \ - --hash=sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049 \ - --hash=sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d \ - --hash=sha256:6c7ecb74c4bd71db68a6bea1edf8da8c34f3d9fe218f038814fd1d310ac76c90 \ - --hash=sha256:7c23c54a00ae43edf48d44066a7ec31e05fdc2eee0be2b8b50dd1903a1db94bb \ - --hash=sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465 \ - --hash=sha256:88c982aac7cb1cbe8cbb4e7f253072b1df872701fcaf48d84ffbb433b6568f24 \ - --hash=sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453 \ - --hash=sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56 \ - --hash=sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48 \ - --hash=sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff \ - --hash=sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460 \ - --hash=sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac \ - --hash=sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900 \ - --hash=sha256:a9b61c19040397970d18d7737375cffd83b1f36a11dd4ad19f83a016f736c3ef \ - --hash=sha256:b4b801ebe0b477be666696bda493a9be8356f1f0057a57f1e35cd26928823e5a \ - --hash=sha256:b95e97e470fe60ed493fd9ae3911d8da4ebac16bd21f87ffa2b7c588bf22ea2c \ - --hash=sha256:bc11d7e8c44a65115d05e2ab9989d1e045125d7be8e05a071a48bc76eb6d6040 \ - --hash=sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9 \ - --hash=sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7 \ - --hash=sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6 \ - --hash=sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b \ - --hash=sha256:d81fdb088defa30eb37bf390bb7dde35d3a83ec112ac8e33d75ab28cc29dd8b0 \ - --hash=sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328 +ml-dtypes==0.6.0 \ + --hash=sha256:008382aeab529df5d3f00501ad9a7dcd64494d4b5b1971fc4c79019e6c1f5010 \ + --hash=sha256:03ce583adfce34ad33aa9e1fc7a8344dcf90ea776cc4ef0e5a48d4eae84e5d20 \ + --hash=sha256:084dfe51a7ad58b171f05115f8226ed4233a454a1611371947e806e76f0c638d \ + --hash=sha256:26b1f1fa4f0435a2946859823f6e2bf06796f1e9f10f5a05b08a5e3c8f46ff69 \ + --hash=sha256:28d676428b104bb9717b0928bc5c5129f2d6b51b6727587cc4289e7bf8713cb5 \ + --hash=sha256:2a3e9d53925597fbffafd2a37048dadeddd0bdaba58058f6ae0869ed709a184d \ + --hash=sha256:3035518e3e19add1a4cac9236ab22888b208a4074912514313ccb2d6d242cde8 \ + --hash=sha256:317be9967fb84b0ce4e80e6b1bf71213d21971621cf6f1e501a63602a95297bf \ + --hash=sha256:31f1ce979d31a357e95aa81812f20412c8c954fa43c44ee3ead1e1c8a78575ef \ + --hash=sha256:37da32aa97749251025666d62372775019594577b9c9e9cfda83bed48d778fdb \ + --hash=sha256:3b4a480aa8fd54a1805b8ac10f3f91763926a74f73c0c364c10f9231854f4170 \ + --hash=sha256:3be9911d953f97cddded4b9961d7b650473b7e55806d20f6176f8356dfe7b38e \ + --hash=sha256:3e169214e0d80ff1c038e1b3017e33c23e43bdf948d42d31de8283111c7e2fa3 \ + --hash=sha256:488c99ab181a2f59d9ec3b12c5fa11ec904e92be2c4ba18cded54dd7501208fe \ + --hash=sha256:5359c588cc62de6f78d7430f06b65853d884955494d86d6ad90b6dd64a3f3a08 \ + --hash=sha256:573b11f3c327e17ef3826d266e676cf1149a1f3016f822a05f2306c55d8246bf \ + --hash=sha256:57ed0d6b4ac5e7868361303a9c57fbcf63b768236ee14456f585dfcf260d0292 \ + --hash=sha256:5a519c9e95a216fbcb8e759793ef7fb40793fc803ed839142d6dc5be9be5bc89 \ + --hash=sha256:5e60251d32ced5598972e4d5e06a2f044341f9291402551a3f6f0ec44f9299b0 \ + --hash=sha256:6c8e39b53e90afda8ce52859c93de4dba3e02b76d85dcf091cc469f9184c6dae \ + --hash=sha256:6eaed129a4afe90694b8685e2f9b6294849f5eda4af9a15be83a4326eeebd775 \ + --hash=sha256:6ec0d244a5bba12239025389ad88bbfb45f9f10e25ab4f678e9a4768ebd47532 \ + --hash=sha256:7728c0420ec1c338564fc8b01015ff2d58567e70f17fedce5a0a7c0308c0d5b9 \ + --hash=sha256:84fa136b8602c8c39e3b6cb24918960cd6f36cade7a70376f56770729cd56510 \ + --hash=sha256:8f490c003369ce60e514a0c3b12374f05274c101fee1bead6740ec8a564032b0 \ + --hash=sha256:9c6ad60af4102789a5c09824004beade2f7f28cd1cd581ee5c170d9dc2fbb00e \ + --hash=sha256:b1b503864fada3f74fabf8d9fee7b4c1cbe956301e6fdece975d5f77c2fce958 \ + --hash=sha256:b76fa1d3f92967d58289ac47ab7458ede66e6f3527fff3e59142aee57d9307cd \ + --hash=sha256:bad8d1dd5bed060a29332b99d63d0e5c2969081e1c6ea54adfbccfdfa783be44 \ + --hash=sha256:ce7563e0b1a4482cbc1b4a6272145e54e4489e54fe7428f94908c3d87103abfa \ + --hash=sha256:d4f1b9329a251e4affe3bb58f4d3e2db22a714396fd7ffb40d0b5db423c24d17 \ + --hash=sha256:d574c2b28921dc72e869df248f1a278f6eee176a1f237c8642e1a71eb15f3977 \ + --hash=sha256:de9d14748dbf3968951436ef514a29c9d1fe438aa680d110134ee2f7a9f9df18 \ + --hash=sha256:e25bb3b0ad1217b60626e4ed45b10ca170c41d99fbe44a12bebc1e07ec4aad55 \ + --hash=sha256:e2d6149f3a57f405bcad5fb41e03218b8373936253f23e1ca84c0108abbc3392 \ + --hash=sha256:e74266ca8e97874a937b7646378c178025650a236584f7474d10d8086a6edea3 \ + --hash=sha256:f4adb4af61516510d786cf8c01851a66f6d3ddfa79e1144deaa5b40d8507231e \ + --hash=sha256:f4f59f83c82ab480e924b988e7b1b4eb4de836dfcf5390c6f59148d1a00e1d02 \ + --hash=sha256:f6cb525101b6b903779188c1e9e9490c343b455ab822883e02cf01e5547338d2 \ + --hash=sha256:fb87f46b4f7ad7b5d3ad8f4b452b024bd4229d44c8ff934798c1fe656210387a From a710f7ce033428c342c9e3912a0131cce594d910 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:46:33 -0700 Subject: [PATCH 08/23] PR #47910: Bump astral-sh/setup-uv from 10.0.0 to 10.0.1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47910 Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 10.0.0 to 10.0.1.
Release notes

Sourced from astral-sh/setup-uv's releases.

v10.0.1 ๐ŸŒˆ Tolerate transient manifest timeouts

Changes

Thank you @โ€‹arguile- for making this action more resilient.

๐Ÿ› Bug fixes

๐Ÿงฐ Maintenance

๐Ÿ“š Documentation

Commits

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=astral-sh/setup-uv&package-manager=github_actions&previous-version=10.0.0&new-version=10.0.1)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Copybara import of the project: -- 466d0bbcf6e70095833a34245cc0a2d78e34aaf8 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>: Bump astral-sh/setup-uv from 10.0.0 to 10.0.1 Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 10.0.0 to 10.0.1. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](https://github.com/astral-sh/setup-uv/compare/ae62891fec2bb8e7d6c99fc78c9fec3a63790f8d...20cfd1bf945f4377ade1205e4dbc17946fc9a30d) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 10.0.1 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Merging this change closes #47910 PiperOrigin-RevId: 972433551 --- third_party/xla/.github/workflows/clang_format.yml | 2 +- third_party/xla/.github/workflows/zizmor.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/xla/.github/workflows/clang_format.yml b/third_party/xla/.github/workflows/clang_format.yml index 477e70220a3bba..359d5b78d2a2d7 100644 --- a/third_party/xla/.github/workflows/clang_format.yml +++ b/third_party/xla/.github/workflows/clang_format.yml @@ -41,7 +41,7 @@ jobs: - name: "Fetch HEAD of main branch" run: git fetch origin main --depth=1 - name: Install uv dependencies - uses: astral-sh/setup-uv@ae62891fec2bb8e7d6c99fc78c9fec3a63790f8d # v10.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 - name: "Run clang-format check" env: TARGET_REF: "origin/${{ github.base_ref || 'main' }}" diff --git a/third_party/xla/.github/workflows/zizmor.yml b/third_party/xla/.github/workflows/zizmor.yml index 9268cac0007cf7..3bcae86a2acdb9 100644 --- a/third_party/xla/.github/workflows/zizmor.yml +++ b/third_party/xla/.github/workflows/zizmor.yml @@ -33,7 +33,7 @@ jobs: persist-credentials: false - name: Install uv dependencies - uses: astral-sh/setup-uv@ae62891fec2bb8e7d6c99fc78c9fec3a63790f8d # v10.0.0 + uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 - name: Run zizmor run: uvx zizmor --format=github .github xla From 97248424c70c57310ef8b7d97c2c31c4099ccc4d Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Fri, 28 Aug 2026 01:47:46 -0700 Subject: [PATCH 09/23] PR #47951: [ROCm] Enable command buffer profiling test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47951 ๐Ÿ“ Summary of Changes - Resolve the active GPU platform instead of hardcoding CUDA/H100. - Build `GpuTargetConfig` from the active StreamExecutor. - Use `xla_test(backends = ["gpu"])` so CUDA and ROCm receive the correct generated targets and CI tags. - Preserve existing CUDA behavior. ๐ŸŽฏ Justification Add more unit test coverage on ROCm platform. ๐Ÿš€ Kind of Contribution ๐Ÿงช Tests Copybara import of the project: -- b20e115b198d8cadec00583b3546a78fa8877bd0 by Lin Chen1 : Enable command buffer profiling test on ROCm Merging this change closes #47951 PiperOrigin-RevId: 972434151 --- third_party/xla/xla/backends/gpu/tests/BUILD | 13 ++++----- .../tests/command_buffer_profiling_test.cc | 27 +++++++++++-------- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 76576adcd14f38..6fbd1a47698d7c 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -1719,13 +1719,10 @@ xla_cc_test( ], ) -xla_cc_test( +xla_test( name = "command_buffer_profiling_test", srcs = ["command_buffer_profiling_test.cc"], - tags = [ - "cuda-only", - "gpu", - ], + backends = ["gpu"], deps = [ "//xla:xla_proto_cc", "//xla/backends/gpu/runtime:command_buffer_thunk", @@ -1736,11 +1733,11 @@ xla_cc_test( "//xla/service:compiled_module", "//xla/service:compiler", "//xla/service:executable", - "//xla/service:gpu_plugin", "//xla/service:gpu_topology", + "//xla/service:platform_util", "//xla/service/gpu:gpu_executable", - "//xla/stream_executor:device_description_proto_cc", - "//xla/stream_executor/cuda:cuda_platform_id", + "//xla/stream_executor:platform", + "//xla/stream_executor:stream_executor_h", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", ], diff --git a/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc b/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc index 303d791f1d706f..791999a6c8b0cb 100644 --- a/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/command_buffer_profiling_test.cc @@ -30,8 +30,9 @@ limitations under the License. #include "xla/service/executable.h" #include "xla/service/gpu/gpu_executable.h" #include "xla/service/gpu_topology.h" -#include "xla/stream_executor/cuda/cuda_platform_id.h" -#include "xla/stream_executor/device_description.pb.h" +#include "xla/service/platform_util.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/stream_executor.h" #include "xla/xla.pb.h" namespace xla::gpu { @@ -57,15 +58,19 @@ TEST_P(CommandBufferProfilingTest, ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, ParseAndReturnVerifiedModule(hlo_text)); - ASSERT_OK_AND_ASSIGN( - stream_executor::GpuTargetConfigProto gpu_target_config_proto, - GetGpuTargetConfig(GpuModel::H100_PCIE)); - ASSERT_OK_AND_ASSIGN( - gpu::GpuTargetConfig gpu_target_config, - gpu::GpuTargetConfig::FromProto(gpu_target_config_proto)); - ASSERT_OK_AND_ASSIGN( - std::unique_ptr compiler, - Compiler::GetForPlatform(stream_executor::cuda::kCudaPlatformId)); + ASSERT_OK_AND_ASSIGN(stream_executor::Platform * platform, + PlatformUtil::GetPlatform("gpu")); + + ASSERT_OK_AND_ASSIGN(std::vector executors, + PlatformUtil::GetStreamExecutors(platform)); + + ASSERT_FALSE(executors.empty()); + + stream_executor::StreamExecutor* executor = executors.front(); + GpuTargetConfig gpu_target_config(executor); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr compiler, + Compiler::GetForPlatform(platform->id())); AotCompilationOptions aot_options(compiler->PlatformId()); aot_options.set_gpu_topology( From 658c58ac916c6272ff6442c1251176c1d288c8d5 Mon Sep 17 00:00:00 2001 From: Bhavani Subramanian Date: Fri, 28 Aug 2026 01:48:36 -0700 Subject: [PATCH 10/23] PR #48007: [XLA:GPU][oneAPI] Fix SYCL heap corruption at process exit Imported from GitHub PR https://github.com/openxla/xla/pull/48007 oneAPI 2026.x tears down libsycl/Level Zero earlier at process exit than pre-2026.x. If SYCL runtime statics such as `device_pool_`, `stream_pool_map_`, and the cached `::sycl::context` run their destructors as usual, they touch already torn-down libsycl/Level Zero state, corrupting the heap (observed as `double free or corruption (!prev)` crashes immediately after tests report PASSED). This PR wraps these statics in `absl::NoDestructor` so they're intentionally leaked instead of destructed at exit. The OS reclaims this memory when the process exits, which is acceptable since this is a one-time allocation per program run. Copybara import of the project: -- 06e491e4cda93ea84b27d48956b1e41d20a839f6 by Bhavani Subramanian : Leak SYCL runtime statics to fix double-free at process exit. oneAPI 2026.x tears down libsycl/Level Zero earlier at process exit than pre-2026.x. If device_pool_, stream_pool_map_, and the cached ::sycl::context run their destructors as usual, they touch already torn-down libsycl/Level Zero state, corrupting the heap ("double free or corruption (!prev)"). Wrap these statics in absl::NoDestructor so they're intentionally leaked instead of destructed at exit. The OS reclaims this memory when the process exits, which is acceptable since this is a one-time allocation per program run. Merging this change closes #48007 PiperOrigin-RevId: 972434558 --- .../xla/xla/stream_executor/sycl/BUILD | 1 + .../xla/stream_executor/sycl/sycl_context.cc | 2 +- .../stream_executor/sycl/sycl_gpu_runtime.cc | 45 ++++++++++--------- .../stream_executor/sycl/sycl_gpu_runtime.h | 16 ++++--- .../sycl/sycl_gpu_runtime_test.cc | 4 +- 5 files changed, 39 insertions(+), 29 deletions(-) diff --git a/third_party/xla/xla/stream_executor/sycl/BUILD b/third_party/xla/xla/stream_executor/sycl/BUILD index 4c22875e53d4ad..0511a4e3894db6 100644 --- a/third_party/xla/xla/stream_executor/sycl/BUILD +++ b/third_party/xla/xla/stream_executor/sycl/BUILD @@ -456,6 +456,7 @@ sycl_library( "//xla/tsl/platform:statusor", "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_context.cc b/third_party/xla/xla/stream_executor/sycl/sycl_context.cc index aed9bd0152bb73..6210396b386483 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_context.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_context.cc @@ -21,7 +21,7 @@ namespace stream_executor::sycl { absl::StatusOr> SyclContext::Create( int device_ordinal) { - ABSL_ASSIGN_OR_RETURN(::sycl::context sycl_context, + ABSL_ASSIGN_OR_RETURN(const ::sycl::context& sycl_context, SyclDevicePool::GetDeviceContext()); return std::make_unique(sycl_context, device_ordinal); } 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 81743e08de4a8b..b10b9e6a852d5b 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 @@ -186,7 +186,7 @@ bool SyclIsHostMemoryRegistered(const ::sycl::device& device, } } -DevicePool SyclDevicePool::device_pool_; +absl::NoDestructor SyclDevicePool::device_pool_; absl::Status SyclDevicePool::InitDevicePool() { static absl::once_flag device_init_flag; @@ -217,29 +217,32 @@ absl::Status SyclDevicePool::InitDevicePool() { "backend. Check oneAPI installation and environment variables."); return; } - device_pool_ = std::move(devices); + // absl::NoDestructor default-constructs an empty DevicePool, so this + // dereference is safe. + *device_pool_ = std::move(devices); }); return init_status; } -absl::StatusOr<::sycl::context> SyclDevicePool::GetDeviceContext() { +absl::StatusOr SyclDevicePool::GetDeviceContext() { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); - static ::sycl::context device_context(device_pool_); - return device_context; + // Leaked for the same reason as SyclDevicePool::device_pool_. + static absl::NoDestructor<::sycl::context> device_context(*device_pool_); + return *device_context; } absl::StatusOr SyclDevicePool::GetDeviceCount() { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); // Cast to int since device_ordinal is usually an int. - return static_cast(device_pool_.size()); + return static_cast(device_pool_->size()); } absl::StatusOr SyclDevicePool::GetDeviceOrdinal( const ::sycl::device& device) { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); - auto it = std::find(device_pool_.begin(), device_pool_.end(), device); - if (it != device_pool_.end()) { - return static_cast(it - device_pool_.begin()); + auto it = std::find(device_pool_->begin(), device_pool_->end(), device); + if (it != device_pool_->end()) { + return static_cast(it - device_pool_->begin()); } return absl::InternalError( "SyclDevicePool::GetDeviceOrdinal failed, got invalid device"); @@ -249,10 +252,10 @@ absl::StatusOr<::sycl::device> SyclDevicePool::GetDevice(int device_ordinal) { ABSL_RETURN_IF_ERROR(SyclDevicePool::InitDevicePool()); ABSL_RETURN_IF_ERROR( IsValidDeviceOrdinal(device_ordinal, "SyclDevicePool::GetDevice")); - return device_pool_[device_ordinal]; + return (*device_pool_)[device_ordinal]; } -StreamPoolMap SyclStreamPool::stream_pool_map_; +absl::NoDestructor SyclStreamPool::stream_pool_map_; absl::Mutex SyclStreamPool::stream_pool_mu_(absl::kConstInit); void SyclAsyncHandler(::sycl::exception_list ex_list) { @@ -269,10 +272,10 @@ void SyclAsyncHandler(::sycl::exception_list ex_list) { absl::StatusOr SyclStreamPool::InitStreamPool(int device_ordinal) { { absl::ReaderMutexLock read_lock(&stream_pool_mu_); - auto it = stream_pool_map_.find(device_ordinal); + auto it = stream_pool_map_->find(device_ordinal); // Returns the existing non-empty stream pool for this device, if available. // The pool may be empty if DestroyStream was called on the last stream. - if (it != stream_pool_map_.end() && !it->second.empty()) { + if (it != stream_pool_map_->end() && !it->second.empty()) { VLOG(2) << "Check 1: Returning existing stream pool for device ordinal " << device_ordinal << " whose size is " << it->second.size(); return &(it->second); @@ -283,14 +286,14 @@ absl::StatusOr SyclStreamPool::InitStreamPool(int device_ordinal) { ::sycl::property::queue::in_order()}; ABSL_ASSIGN_OR_RETURN(::sycl::device sycl_device, SyclDevicePool::GetDevice(device_ordinal)); - ABSL_ASSIGN_OR_RETURN(::sycl::context sycl_context, + ABSL_ASSIGN_OR_RETURN(const ::sycl::context& sycl_context, SyclDevicePool::GetDeviceContext()); VLOG(2) << "Creating new stream pool for device ordinal " << device_ordinal; absl::MutexLock write_lock(&stream_pool_mu_); - auto it = stream_pool_map_.find(device_ordinal); + auto it = stream_pool_map_->find(device_ordinal); // Double-checks that another thread has not already created the pool. - if (it != stream_pool_map_.end() && !it->second.empty()) { + if (it != stream_pool_map_->end() && !it->second.empty()) { VLOG(2) << "Check 2: Returning existing stream pool for device ordinal " << device_ordinal << " whose size is " << it->second.size(); return &(it->second); @@ -301,9 +304,9 @@ absl::StatusOr SyclStreamPool::InitStreamPool(int device_ordinal) { // Use assignment (not insert) to update the stream pool if it was // previously destroyed. - stream_pool_map_[device_ordinal] = std::move(stream_pool); + (*stream_pool_map_)[device_ordinal] = std::move(stream_pool); - return &(stream_pool_map_[device_ordinal]); + return &((*stream_pool_map_)[device_ordinal]); } absl::StatusOr SyclStreamPool::GetDefaultStream(int device_ordinal) { @@ -344,7 +347,7 @@ absl::StatusOr SyclStreamPool::GetOrCreateStream( ::sycl::property::queue::in_order()}; ABSL_ASSIGN_OR_RETURN(::sycl::device sycl_device, SyclDevicePool::GetDevice(device_ordinal)); - ABSL_ASSIGN_OR_RETURN(::sycl::context sycl_context, + ABSL_ASSIGN_OR_RETURN(const ::sycl::context& sycl_context, SyclDevicePool::GetDeviceContext()); stream_pool->push_back(std::make_shared<::sycl::queue>( sycl_context, sycl_device, SyclAsyncHandler, prop_list)); @@ -406,7 +409,7 @@ absl::Status SyclStreamPool::DestroyStream(int device_ordinal, void SyclStreamPool::Reset() { absl::MutexLock write_lock(&stream_pool_mu_); - for (auto& [device_ordinal, stream_pool] : stream_pool_map_) { + for (auto& [device_ordinal, stream_pool] : *stream_pool_map_) { for (auto& stream_handle : stream_pool) { if (stream_handle) { stream_handle->wait(); @@ -415,7 +418,7 @@ void SyclStreamPool::Reset() { } stream_pool.clear(); } - stream_pool_map_.clear(); + stream_pool_map_->clear(); } absl::StatusOr SyclGetTimerProperties(int device_ordinal) { diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h index 0c3fcc76e3724b..ce0814957f7e04 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h +++ b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime.h @@ -26,6 +26,7 @@ limitations under the License. #include #include "absl/base/attributes.h" +#include "absl/base/no_destructor.h" #include "absl/container/flat_hash_map.h" #include "absl/strings/ascii.h" #include "xla/stream_executor/sycl/sycl_status.h" @@ -53,7 +54,7 @@ class SyclDevicePool { // This function assumes that the device pool is not modified after // initialization. If this assumption is violated, the context may become // invalid. - static absl::StatusOr<::sycl::context> GetDeviceContext(); + static absl::StatusOr GetDeviceContext(); // Returns the number of devices in the pool. static absl::StatusOr GetDeviceCount(); @@ -65,8 +66,11 @@ class SyclDevicePool { static absl::StatusOr<::sycl::device> GetDevice(int device_ordinal); private: - // The underlying device pool. - static DevicePool device_pool_; + // Leaked: libsycl/Level Zero teardown order relative to this static is + // oneAPI-version-dependent (after it pre-2026.x, before it in 2026.x+), + // making its destructor unsafe to run, so cleanup is deferred to the OS + // instead. Acceptable since this is a one-time allocation per program run. + static absl::NoDestructor device_pool_; // Thread-safe initialization of device_pool_ with all Level-Zero backend GPUs // using absl::call_once. @@ -120,8 +124,10 @@ class SyclStreamPool { static absl::Mutex stream_pool_mu_; // The underlying stream pool for each device. The device ordinal - // is used as the key. - static StreamPoolMap stream_pool_map_ ABSL_GUARDED_BY(stream_pool_mu_); + // is used as the key. Leaked for the same reason as + // SyclDevicePool::device_pool_. + static absl::NoDestructor stream_pool_map_ + ABSL_GUARDED_BY(stream_pool_mu_); // Initializes and returns a pointer to the stream pool for the given device // ordinal. diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc index 1e733a0d18edf4..190a4141539e8a 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_gpu_runtime_test.cc @@ -126,9 +126,9 @@ TEST_F(SyclGpuRuntimeTest, GetDeviceOrdinal) { TEST_F(SyclGpuRuntimeTest, TestStaticDeviceContext) { // Verify that GetDeviceContext returns the same context instance on multiple // calls. - TF_ASSERT_OK_AND_ASSIGN(::sycl::context saved_sycl_context, + TF_ASSERT_OK_AND_ASSIGN(const ::sycl::context& saved_sycl_context, SyclDevicePool::GetDeviceContext()); - TF_ASSERT_OK_AND_ASSIGN(::sycl::context current_sycl_context, + TF_ASSERT_OK_AND_ASSIGN(const ::sycl::context& current_sycl_context, SyclDevicePool::GetDeviceContext()); EXPECT_EQ(saved_sycl_context, current_sycl_context); } From 216232f7bd23e265ec315d3175acdfbb4e30660a Mon Sep 17 00:00:00 2001 From: linchen1-robot Date: Fri, 28 Aug 2026 01:50:55 -0700 Subject: [PATCH 11/23] PR #47950: [ROCm] Enable dynamic slice fusion rewriter v2 test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47950 ๐Ÿ“ Summary of Changes - Remove the stale `cuda-only` tag from `dynamic_slice_fusion_rewriter_v2_test`. - Enable this backend-agnostic test in ROCm CI. - Preserve existing CUDA behavior. ๐ŸŽฏ Justification Add more unit test coverage on ROCm platform ๐Ÿš€ Kind of Contribution ๐Ÿงช Tests Copybara import of the project: -- fbac1d2bdc2d6021168e5a484fcf49f6b3a6f63f by Lin Chen1 : Enable dynamic slice fusion rewriter v2 test on ROCm Merging this change closes #47950 PiperOrigin-RevId: 972435961 --- third_party/xla/xla/backends/gpu/transforms/BUILD | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 0782615d4ddab4..5d521b072e4bc5 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -1734,10 +1734,7 @@ cc_library( xla_cc_test( name = "dynamic_slice_fusion_rewriter_v2_test", srcs = ["dynamic_slice_fusion_rewriter_v2_test.cc"], - tags = [ - "cuda-only", - "gpu", - ], + tags = ["gpu"], deps = [ ":dynamic_slice_annotator", ":dynamic_slice_fusion", From ce63e2b635d17527d542f92deb513ec5bd75260e Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:51:01 -0700 Subject: [PATCH 12/23] PR #47909: Bump github/codeql-action/upload-sarif from 4.37.6 to 4.37.7 Imported from GitHub PR https://github.com/openxla/xla/pull/47909 Bumps [github/codeql-action/upload-sarif](https://github.com/github/codeql-action) from 4.37.6 to 4.37.7.
Release notes

Sourced from github/codeql-action/upload-sarif's releases.

v4.37.7

  • Update default CodeQL bundle version to 2.26.3. #4085
Changelog

Sourced from github/codeql-action/upload-sarif's changelog.

CodeQL Action Changelog

See the releases page for the relevant changes to the CodeQL CLI and language packs.

[UNRELEASED]

  • Update default CodeQL bundle version to 2.26.4. #4106

4.37.8 - 21 Aug 2026

No user facing changes.

4.37.7 - 13 Aug 2026

  • Update default CodeQL bundle version to 2.26.3. #4085

4.37.6 - 04 Aug 2026

  • Changed the default filepath for the new remote file address format that was introduced in CodeQL Action 4.37.0 / 3.37.0 to .github/codeql-config.yml to align it with the suggested path that is used elsewhere. #4070

4.37.5 - 03 Aug 2026

  • Fixed a bug where a network error while streaming the download of the CodeQL bundle could terminate the init Action instead of falling back to downloading the bundle before extracting it. #4061

4.37.4 - 29 Jul 2026

  • This version of the CodeQL Action adds support for the tools input for the codeql-action/init step to be specified using a github-codeql-tools repository property. This feature will gradually be rolled out following the release of this version. Once rolled out, this allows for the CodeQL CLI version that is used in GitHub-managed workflows, such as Default Setup, to be set to a custom value. For example, customers who run into issues with rate limits when a new CodeQL CLI version is released can set the value to toolcache to always use the CodeQL CLI version that is available in the runner toolcache. For Advanced Setup workflows, the value provided for tools in the workflow definition always takes precedence unless the value of the repository property starts with !. #4037
  • Update default CodeQL bundle version to 2.26.2. #4051

4.37.3 - 22 Jul 2026

No user facing changes.

4.37.2 - 21 Jul 2026

  • The new address format for the config-file input that was introduced in CodeQL Action 4.37.0 is now enabled by default. In addition to the format described there, the remote= prefix can now be used to explicitly indicate that the input refers to a remote file. All previous input formats continue to be accepted as well. #4023
  • The CodeQL Action can now make use of configured private registries in Default Setup to retrieve CodeQL configuration files from remote repositories that require authentication. This will allow customers to store their CodeQL configuration in a single repository that can then be referenced by Default Setup workflows in other repositories. We expect to roll this and other, related changes out to everyone in July. #4007

4.37.1 - 16 Jul 2026

  • Upcoming breaking change: Add a deprecation warning for customers using CodeQL version 2.20.6 and earlier. These versions of CodeQL were discontinued on 1 July 2026 alongside GitHub Enterprise Server 3.16, and will be unsupported by the next minor release of the CodeQL Action. #3956
  • Update default CodeQL bundle version to 2.26.1. #4019

4.37.0 - 08 Jul 2026

  • Update default CodeQL bundle version to 2.26.0. #3995
  • In addition to the existing input format, the config-file input for the codeql-action/init step will soon support a new [owner/]repo[@ref][:path] format. All components except the repository name are optional. If omitted, owner defaults to the same owner as the repository the analysis is running for, ref to main, and path to .github/codeql-action.yaml. Support for this format ships in this version of the CodeQL Action, but will only be enabled over the coming weeks. #3973

4.36.3 - 01 Jul 2026

... (truncated)

Commits
  • ff2f1c6 Merge pull request #4093 from github/update-v4.37.7-be7a3dbb8
  • 951a133 Update changelog for v4.37.7
  • be7a3db Merge pull request #4087 from github/dependabot/npm_and_yarn/npm-minor-0aa561...
  • 9310334 Merge pull request #4086 from github/mbg/thread-action-state-to-codeql
  • b4d8a54 Rebuild
  • ab5db25 Bump the npm-minor group across 1 directory with 8 updates
  • 38055a3 Drop logger from databaseInitCluster in interface
  • 1f87aed Merge pull request #4085 from github/update-bundle/codeql-bundle-v2.26.3
  • dc1b98a Make logger available to getCodeQLForCmd
  • 6f0220e Merge pull request #4084 from github/navntoft/bump-undici
  • Additional commits viewable in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=github/codeql-action/upload-sarif&package-manager=github_actions&previous-version=4.37.6&new-version=4.37.7)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Copybara import of the project: -- b42df13028500dd6dc74cead0ea0a2e81d1f9f66 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>: Bump github/codeql-action/upload-sarif from 4.37.6 to 4.37.7 Bumps [github/codeql-action/upload-sarif](https://github.com/github/codeql-action) from 4.37.6 to 4.37.7. - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/github/codeql-action/compare/5595ccaf912efad79be6eef63a5619ff05969be3...ff2f1c621b7f889edc0d3c761ac2e6a3f8cdb0dd) --- updated-dependencies: - dependency-name: github/codeql-action/upload-sarif dependency-version: 4.37.7 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Merging this change closes #47909 PiperOrigin-RevId: 972436016 --- third_party/xla/.github/workflows/scorecards-analysis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/xla/.github/workflows/scorecards-analysis.yml b/third_party/xla/.github/workflows/scorecards-analysis.yml index 81077b90ea3b6f..f08174a83b51d0 100644 --- a/third_party/xla/.github/workflows/scorecards-analysis.yml +++ b/third_party/xla/.github/workflows/scorecards-analysis.yml @@ -67,6 +67,6 @@ jobs: # Upload the results to GitHub's code scanning dashboard (optional). # Commenting out will disable upload of results to your repo's Code Scanning dashboard - name: "Upload to code-scanning" - uses: github/codeql-action/upload-sarif@5595ccaf912efad79be6eef63a5619ff05969be3 # v4.37.6 + uses: github/codeql-action/upload-sarif@ff2f1c621b7f889edc0d3c761ac2e6a3f8cdb0dd # v4.37.7 with: sarif_file: results.sarif From 630ae1a9dd0220e8fb961aedd6542bf81cc1725e Mon Sep 17 00:00:00 2001 From: Shyamli Agrawal Date: Fri, 28 Aug 2026 02:19:31 -0700 Subject: [PATCH 13/23] Stop uploading AutotuneResults with optimized HLO. PiperOrigin-RevId: 972449593 --- .../xla/xla/service/gpu/gpu_compiler.cc | 5 +---- third_party/xla/xla/tools/xla_compile_lib.cc | 1 + .../xla/xla/tools/xla_gpu_compile_lib_test.cc | 19 +++++++++++++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 4ca8fef163aebd..a5408c0f2f3587 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -2379,16 +2379,13 @@ absl::StatusOr> GpuCompiler::RunHloPasses( DumpHloModuleMetadataIfEnabled(module.get()); - AutotuneResults autotune_results; if (stream_exec != nullptr) { - ABSL_RETURN_IF_ERROR( - AutotunerCache::SerializeAutotuneResults(&autotune_results)); ABSL_RETURN_IF_ERROR(SerializeAutotuneResultsToFile(debug_opts)); } std::optional optimized_fingerprint; if (should_upload_hlo_modules) { optimized_fingerprint = - MaybeUploadOptimizedGpuSymbols(module.get(), autotune_results); + MaybeUploadOptimizedGpuSymbols(module.get(), AutotuneResults()); } if (unoptimized_fingerprint.has_value() && optimized_fingerprint.has_value()) { diff --git a/third_party/xla/xla/tools/xla_compile_lib.cc b/third_party/xla/xla/tools/xla_compile_lib.cc index 176f4819ecd522..8f5a59d0e96aa4 100644 --- a/third_party/xla/xla/tools/xla_compile_lib.cc +++ b/third_party/xla/xla/tools/xla_compile_lib.cc @@ -333,6 +333,7 @@ absl::StatusOr LoadAutotuneDataFromModule(HloModuleAndMetadata* mod, if (auto* data = static_cast( mod->backend_specific_data.get()); data != nullptr && data->autotune_results.has_value() && + !data->autotune_results->results().empty() && mod->hlo_module->config().debug_options().xla_gpu_autotune_level() > 0) { ABSL_RETURN_IF_ERROR( diff --git a/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc b/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc index 03ce87a90a5eff..c850c02df1b3f9 100644 --- a/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc +++ b/third_party/xla/xla/tools/xla_gpu_compile_lib_test.cc @@ -263,6 +263,25 @@ TEST_F(XlaCompileLibTest, EXPECT_TRUE(gpu::AutotunerCache::ResultCacheIsEmpty()); } +TEST_F(XlaCompileLibTest, + LoadAutotuneDataGpuDataPresentEmptyAndAutotuningEnabled) { + gpu::AutotunerCache::ClearAutotuneResults(); + + HloModuleAndMetadata mod; + mod.hlo_module = std::move(module_); + auto data = std::make_unique(); + data->autotune_results.emplace(); // empty results + mod.backend_specific_data = std::move(data); + + DebugOptions opts = mod.hlo_module->config().debug_options(); + opts.set_xla_gpu_autotune_level(3); + mod.hlo_module->mutable_config().set_debug_options(opts); + + EXPECT_THAT(internal::LoadAutotuneDataFromModule(&mod, BackendType::kGpu), + absl_testing::IsOkAndHolds(false)); + EXPECT_TRUE(gpu::AutotunerCache::ResultCacheIsEmpty()); +} + TEST_F(XlaCompileLibTest, MainForGpuForceAutoLayout) { static constexpr absl::string_view kHloText = R"( HloModule f From 33f9a169c961a4fba92a7fdc79dc160ea500667d Mon Sep 17 00:00:00 2001 From: clarke <149320067+clarkechong@users.noreply.github.com> Date: Fri, 28 Aug 2026 02:21:31 -0700 Subject: [PATCH 14/23] PR #47917: [ROCm: XLA Profiler] Write VGPR and LDS kernel stats to kernel_info XStat for XProf MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47917 ## ๐Ÿ“ Summary of Changes - Populate ROCm `kernel_details` XStat with: - Architecture + accumulator VGPRs per work-item. - Static LDS from kernel symbol metadata. - Dynamic LDS derived from dispatch total minus static LDS. - XProf-compatible `regs`, `static_shared`, and `dynamic_shared` keys. - Adjust inline comments regarding existing `private_segment_size` and `group_segment_size` stats. ## ๐ŸŽฏ Justification On AMD GPU traces, XProf Kernel Stats page reads 0 for "Registers per thread" and "Shared Mem Bytes". This is because XProf parses these fields on specific keys in the `kernel_info` XStat, `regs`, `static_shared` and `dynamic_shared`, none of which we were populating from `rocm_collector`. ### Implementation detail: - Why do we need static/dynamic group segment size separately when we already have total group segment size? - XProf parses on static, dynamic group size separately. The group segment size stat provided by rocprofiler combines static and dynamic, so we need to manually decompose into static and dynamic attribution. - Why not add dynamic mem member (eg. `dynamic_group_segment_size`) to KernelDetails struct if we are already adding `static_group_segment_size`? - The static LDS size is fixed by the compiled kernel symbol and does not vary between dispatches of that symbol. Rocprofiler provides the total LDS allocation at dispatch time. The dynamic size is derived as (dispatch total - static LDS) and can vary between dispatches due to runtime allocation and padding. - Why not just change XProf to parse our `group_mem` instead of following this existing static + dynamic separation? - Would need to vendor gate in XProf as NVIDIA provides static, dynamic separately. Besides, adding separation for static vs dynamic LDS is a useful detail to have. - Rationale for defining registers per work item as arch VGPR + accum VGPR - `registers_per_work_item` sums `arch_vgpr_count + accum_vgpr_count` because both are per-lane allocations consuming the VGPR budget. SGPRs are per-wavefront (not per work item/thread) scalar registers in a separate register file and so are excluded. ### Evidence (XProf v2.23.1, profiled via JAX) VIEWING AMD TRACE GENERATED BEFORE PR CHANGES: image VIEWING AMD TRACE AFTER PR CHANGES: image ## ๐Ÿš€ Kind of Contribution ๐Ÿ› Bug Fix ## ๐Ÿ“Š Benchmark (for Performance Improvements) N/A ## ๐Ÿงช Unit Tests: `TEST(RocmCollectorTest, ToXStatDecomposesDispatchGroupMemory)` Verifies correct decomposition of total group-segment size into static and dynamic LDS, including underflow protection. ## ๐Ÿงช Execution Tests: N/A Copybara import of the project: -- 820a26cc795bab9176983bfce6e8c5e91324ca47 by Clarke Chong : Write VGPR and LDS kernel stats to kernel_info XStat for XProf Merging this change closes #47917 PiperOrigin-RevId: 972450467 --- .../xla/backends/profiler/gpu/rocm_collector.h | 12 ++++++++++-- .../backends/profiler/gpu/rocm_collector_test.cc | 15 +++++++++++++++ .../xla/xla/backends/profiler/gpu/rocm_tracer.cc | 8 +++++++- .../xla/backends/profiler/gpu/rocm_tracer_utils.h | 10 +++++++--- 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h index cafdc0bde36591..602e25f8302eb0 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector.h @@ -50,8 +50,16 @@ inline std::string ToXStat(const KernelDetails& kernel_info, grid_z = kernel_info.workgroup_z != 0 ? kernel_info.grid_z / kernel_info.workgroup_z : 0; - - return absl::StrCat(" grid:", grid_x, ",", grid_y, ",", grid_z, + const uint32_t dynamic_group_segment_size = + kernel_info.group_segment_size > kernel_info.static_group_segment_size + ? kernel_info.group_segment_size - + kernel_info.static_group_segment_size + : 0; + + return absl::StrCat("regs:", kernel_info.registers_per_work_item, + " static_shared:", kernel_info.static_group_segment_size, + " dynamic_shared:", dynamic_group_segment_size, + " grid:", grid_x, ",", grid_y, ",", grid_z, " block:", kernel_info.workgroup_x, ",", kernel_info.workgroup_y, ",", kernel_info.workgroup_z, " private_mem:", kernel_info.private_segment_size, diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc index 51cba12620f24c..65cd33905a1298 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_collector_test.cc @@ -36,6 +36,21 @@ namespace test { using tsl::profiler::FindOrAddMutablePlaneWithName; using tsl::profiler::XSpace; +TEST(RocmCollectorTest, ToXStatDecomposesDispatchGroupMemory) { + KernelDetails kernel_info{}; + kernel_info.group_segment_size = 1536; + kernel_info.static_group_segment_size = 1024; + + EXPECT_NE(ToXStat(kernel_info, /*occupancy_pct=*/0) + .find("static_shared:1024 dynamic_shared:512"), + std::string::npos); + + kernel_info.group_segment_size = 512; + EXPECT_NE(ToXStat(kernel_info, /*occupancy_pct=*/0) + .find("static_shared:1024 dynamic_shared:0"), + std::string::npos); +} + TEST(RocmCollectorTest, TestAddKernelEventAndExport) { RocmTraceCollectorOptions options; options.max_callback_api_events = 100; diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc index f9c37cfac6cd1d..5a47fb709fb7e7 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer.cc @@ -361,7 +361,13 @@ void RocmTracer::KernelEvent(const rocprofiler_record_header_t* hdr, }; auto it = kernel_info_.find(kinfo.kernel_id); - if (it != kernel_info_.end()) trace_event->name = it->second.name; + if (it != kernel_info_.end()) { + trace_event->name = it->second.name; + const auto& sym = it->second.data; + trace_event->kernel_info.registers_per_work_item = + sym.arch_vgpr_count + sym.accum_vgpr_count; + trace_event->kernel_info.static_group_segment_size = sym.group_segment_size; + } } void RocmTracer::EmitMarkerEvent(std::string label, uint64_t start_ns, diff --git a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h index 7c93538f9ce640..0de4064154dfe0 100644 --- a/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h +++ b/third_party/xla/xla/backends/profiler/gpu/rocm_tracer_utils.h @@ -59,11 +59,15 @@ struct MemsetDetails { }; struct KernelDetails { - // The amount of private memory used by kernel, - // number of register per thread (register spillage if > 0) + // Total dispatch-time private-segment (scratch) bytes per work-item. uint32_t private_segment_size; - // The amount of shared memory (SMEM) + // Total dispatch-time group-segment (LDS) bytes per workgroup. Includes + // static and dynamic LDS allocation. uint32_t group_segment_size; + // Architecture and accumulator VGPRs allocated per work-item. + uint32_t registers_per_work_item; + // Static group-segment (LDS) bytes per workgroup from the kernel symbol. + uint32_t static_group_segment_size; // X-dimension of a workgroup (grid.x*block.x) uint32_t workgroup_x; // Y-dimension of a workgroup (grid.x*block.x) From fe07098a3788ec856ec9efc3697514538967a531 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Fri, 28 Aug 2026 02:45:50 -0700 Subject: [PATCH 15/23] Fall back to default computation placer and remove alwayslink ComputationPlacer::GetForPlatform now returns a default ComputationPlacer instance when no platform-specific placer is registered. This removes the need for computation_placer.cc to hardcode static registrations for common platforms (host, cuda, rocm, sycl), eliminates dependencies on platform-specific platform IDs, and allows removing alwayslink = True from the target. PiperOrigin-RevId: 972461341 --- .../pjrt/gpu/se_gpu_topology_description.cc | 5 +-- third_party/xla/xla/service/BUILD | 8 +--- third_party/xla/xla/service/backend.cc | 4 +- .../xla/xla/service/computation_placer.cc | 39 ++++--------------- .../xla/xla/service/computation_placer.h | 7 ++-- .../xla/service/computation_placer_test.cc | 39 +++++++++++++++++++ 6 files changed, 56 insertions(+), 46 deletions(-) diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc index 9a1f98f70f82ef..387b22917dd284 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc @@ -337,9 +337,8 @@ StreamExecutorGpuTopologyDescription::GetDefaultDeviceAssignment( stream_executor::PlatformId se_platform_id, StreamExecutorPlatformIdMapping::Global().GetStreamExecutorPlatformId( platform_id())); - ABSL_ASSIGN_OR_RETURN(auto* placer, - ComputationPlacer::GetForPlatform(se_platform_id)); - return placer->AssignDevices(num_replicas, num_partitions); + return ComputationPlacer::GetForPlatform(se_platform_id) + ->AssignDevices(num_replicas, num_partitions); } absl::StatusOr diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index bab0c3b7b4031a..ec981df50f3178 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -3676,19 +3676,13 @@ cc_library( hdrs = ["computation_placer.h"], deps = [ ":device_assignment", - "//xla:util", "//xla/stream_executor:platform_id", - "//xla/stream_executor/cuda:cuda_platform_id", - "//xla/stream_executor/host:host_platform_id", - "//xla/stream_executor/rocm:rocm_platform_id", - "//xla/stream_executor/sycl:sycl_platform_id", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", ], - alwayslink = True, # Contains per-platform computation placer registration ) xla_cc_test( @@ -3697,6 +3691,8 @@ xla_cc_test( deps = [ ":computation_placer", ":device_assignment", + "//xla/stream_executor:platform_id", + "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest_main", ], ) diff --git a/third_party/xla/xla/service/backend.cc b/third_party/xla/xla/service/backend.cc index 8f3e10fdb268fb..016e295a2316f1 100644 --- a/third_party/xla/xla/service/backend.cc +++ b/third_party/xla/xla/service/backend.cc @@ -190,8 +190,8 @@ CreateGpuAllocators(const se::Platform* platform, PlatformUtil::GetStreamExecutors(platform, options.allowed_devices())); ABSL_ASSIGN_OR_RETURN(auto transfer_manager, TransferManager::GetForPlatform(platform)); - ABSL_ASSIGN_OR_RETURN(auto computation_placer, - ComputationPlacer::GetForPlatform(platform->id())); + ComputationPlacer* computation_placer = + ComputationPlacer::GetForPlatform(platform->id()); std::unique_ptr backend(new Backend( platform, std::move(compiler), stream_executors, transfer_manager, computation_placer, options.intra_op_parallelism_threads())); diff --git a/third_party/xla/xla/service/computation_placer.cc b/third_party/xla/xla/service/computation_placer.cc index db8bce900e5290..82730dfca92136 100644 --- a/third_party/xla/xla/service/computation_placer.cc +++ b/third_party/xla/xla/service/computation_placer.cc @@ -24,12 +24,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "xla/service/device_assignment.h" -#include "xla/stream_executor/cuda/cuda_platform_id.h" -#include "xla/stream_executor/host/host_platform_id.h" #include "xla/stream_executor/platform_id.h" -#include "xla/stream_executor/rocm/rocm_platform_id.h" -#include "xla/stream_executor/sycl/sycl_platform_id.h" -#include "xla/util.h" namespace xla { @@ -62,6 +57,11 @@ PlacerFactoryMap& GetPlatformComputationPlacers() { static PlacerFactoryMap* const r = new PlacerFactoryMap; return *r; } + +ComputationPlacer* GetDefaultComputationPlacer() { + static auto* const default_placer = new ComputationPlacer; + return default_placer; +} } // namespace /* static */ @@ -77,16 +77,14 @@ void ComputationPlacer::RegisterComputationPlacer( } /* static */ -absl::StatusOr ComputationPlacer::GetForPlatform( +ComputationPlacer* ComputationPlacer::GetForPlatform( se::PlatformId platform_id) { absl::MutexLock lock(placer_mutex); PlacerFactoryMap& placers = GetPlatformComputationPlacers(); auto it = placers.find(platform_id); if (it == placers.end()) { - return NotFound( - "Could not find registered computation placer for platform %s", - platform_id->ToName()); + return GetDefaultComputationPlacer(); } PlacerState& state = it->second; @@ -94,28 +92,7 @@ absl::StatusOr ComputationPlacer::GetForPlatform( // Lazily create the computation placer the first time it is needed. state.placer = state.creation_function(); } - return state.placer.get(); + return state.placer ? state.placer.get() : GetDefaultComputationPlacer(); } } // namespace xla - -namespace { -// registering default computation placer factory for common platforms. -std::unique_ptr DefaultComputationPlacer() { - return std::make_unique(); -} - -bool InitModule() { - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::host::kHostPlatformId, DefaultComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::cuda::kCudaPlatformId, DefaultComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::rocm::kROCmPlatformId, DefaultComputationPlacer); - xla::ComputationPlacer::RegisterComputationPlacer( - stream_executor::sycl::kSyclPlatformId, DefaultComputationPlacer); - return true; -} - -bool module_initialized = InitModule(); -} // namespace diff --git a/third_party/xla/xla/service/computation_placer.h b/third_party/xla/xla/service/computation_placer.h index 84ed30aaf49289..2f0bbe33d55679 100644 --- a/third_party/xla/xla/service/computation_placer.h +++ b/third_party/xla/xla/service/computation_placer.h @@ -43,10 +43,9 @@ class ComputationPlacer { static void RegisterComputationPlacer(se::PlatformId platform_id, CreationFunction creation_function); - // Returns the computation placer singleton pointer if it is available for the - // given platform, or an error status if it is not. - static absl::StatusOr GetForPlatform( - se::PlatformId platform_id); + // Returns the computation placer singleton pointer registered for the given + // platform, or the default computation placer if none is registered. + static ComputationPlacer* GetForPlatform(se::PlatformId platform_id); private: ComputationPlacer(const ComputationPlacer&) = delete; diff --git a/third_party/xla/xla/service/computation_placer_test.cc b/third_party/xla/xla/service/computation_placer_test.cc index b30b6b6e39aa85..acce0bc0a02bd5 100644 --- a/third_party/xla/xla/service/computation_placer_test.cc +++ b/third_party/xla/xla/service/computation_placer_test.cc @@ -15,9 +15,16 @@ limitations under the License. #include "xla/service/computation_placer.h" +#include + #include #include +#include "absl/status/statusor.h" #include "xla/service/device_assignment.h" +#include "xla/stream_executor/platform_id.h" + +PLATFORM_DEFINE_ID(kUnregisteredPlatformId, UnregisteredPlatform); +PLATFORM_DEFINE_ID(kCustomPlatformId, CustomPlatform); namespace xla { namespace { @@ -33,5 +40,37 @@ TEST(ComputationPlacerTest, Basic) { EXPECT_EQ(da(0, 1), 4); } +TEST(ComputationPlacerTest, GetForPlatformUnregisteredReturnsDefaultPlacer) { + ComputationPlacer* placer = + ComputationPlacer::GetForPlatform(kUnregisteredPlatformId); + ASSERT_NE(placer, nullptr); + ASSERT_OK_AND_ASSIGN(DeviceAssignment da, placer->AssignDevices(2, 2)); + EXPECT_EQ(da(0, 0), 0); + EXPECT_EQ(da(1, 0), 1); + EXPECT_EQ(da(0, 1), 2); + EXPECT_EQ(da(1, 1), 3); +} + +TEST(ComputationPlacerTest, GetForPlatformRegisteredReturnsCustomPlacer) { + class CustomPlacer : public ComputationPlacer { + public: + absl::StatusOr AssignDevices( + int replica_count, int computation_count) override { + DeviceAssignment assignment(replica_count, computation_count); + assignment.Fill(42); + return assignment; + } + }; + + ComputationPlacer::RegisterComputationPlacer( + kCustomPlatformId, []() { return std::make_unique(); }); + + ComputationPlacer* placer = + ComputationPlacer::GetForPlatform(kCustomPlatformId); + ASSERT_NE(placer, nullptr); + ASSERT_OK_AND_ASSIGN(DeviceAssignment da, placer->AssignDevices(1, 1)); + EXPECT_EQ(da(0, 0), 42); +} + } // namespace } // namespace xla From a33c850682d6e01c5ac0482cf1dea56fdeb21dad Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 02:47:38 -0700 Subject: [PATCH 16/23] [NCCL] Upgrade XLA NCCL version to 2.30.7 PiperOrigin-RevId: 972462168 --- .bazelrc | 4 +- MODULE.bazel | 6 +- .../requirements_lock_3_10.txt | 6 +- .../requirements_lock_3_11.txt | 6 +- .../requirements_lock_3_12.txt | 6 +- .../nvidia-requirements.txt | 4 +- ci/official/utilities/code_check_full.bats | 4 +- requirements_lock_3_10.txt | 14 +- requirements_lock_3_11.txt | 11 +- requirements_lock_3_12.txt | 13 +- requirements_lock_3_13.txt | 13 +- requirements_lock_3_14.txt | 13 +- third_party/nccl/nccl_wheel.patch | 78 +++++---- third_party/nccl/workspace.bzl | 6 +- third_party/xla/tensorflow.bazelrc | 2 +- .../xla/third_party/nccl/archive.BUILD | 3 +- .../xla/third_party/nccl/archive.patch | 161 +++++++++--------- .../xla/third_party/nccl/workspace.bzl | 6 +- 18 files changed, 183 insertions(+), 173 deletions(-) diff --git a/.bazelrc b/.bazelrc index fa9874a91a2302..c41207afe93d54 100644 --- a/.bazelrc +++ b/.bazelrc @@ -293,13 +293,13 @@ common:mkl_aarch64 --config=mkl_aarch64_threadpool common:cuda12_version --repo_env=HERMETIC_CUDA_VERSION="12.5.1" common:cuda12_version --repo_env=HERMETIC_CUDNN_VERSION="9.10.0" common:cuda12_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" -common:cuda12_version --repo_env=HERMETIC_NCCL_VERSION="2.29.7" +common:cuda12_version --repo_env=HERMETIC_NCCL_VERSION="2.30.7" common:cuda12_version --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,sm_89,compute_90" common:cuda13_version --repo_env=HERMETIC_CUDA_VERSION="13.0.0" common:cuda13_version --repo_env=HERMETIC_CUDNN_VERSION="9.12.0" common:cuda13_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.3.20" -common:cuda13_version --repo_env=HERMETIC_NCCL_VERSION="2.29.7" +common:cuda13_version --repo_env=HERMETIC_NCCL_VERSION="2.30.7" common:cuda13_version --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_75,sm_80,sm_90,sm_100,compute_120" common:cuda_version --config=cuda12_version diff --git a/MODULE.bazel b/MODULE.bazel index 6d285fee56ca40..f3177c24804133 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -86,9 +86,9 @@ bazel_dep(name = "libprotobuf-mutator", version = "1.3", repo_name = "libprotobu bazel_dep(name = "rules_ml_toolchain") archive_override( module_name = "rules_ml_toolchain", - integrity = "sha256-OwVoeEJCcEHGXRvC9K7aOjB5VXEg9b6MNGkAh6iMXeU=", - strip_prefix = "rules_ml_toolchain-2eddbc595cc0bbe650c2640204f66b14f015f1a8", - urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/2eddbc595cc0bbe650c2640204f66b14f015f1a8.tar.gz"], + integrity = "sha256-izPVqzV2I7ox7EYzRzcRY6ogeLcHUr4LRP2xMrkWh6A=", + strip_prefix = "rules_ml_toolchain-236a498d8624bef88e35ce2518a833204b0db19f", + urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/236a498d8624bef88e35ce2518a833204b0db19f.tar.gz"], ) bazel_dep(name = "xla", repo_name = "xla") diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt index f98cc872f05cd7..137bc9e338a4e4 100644 --- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt +++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt @@ -519,9 +519,9 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt index 9ddfca552b7545..8b37e5bd8b7bde 100644 --- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt +++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt @@ -519,9 +519,9 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt index 15e097176cdbd4..7c0230a0cbe19b 100644 --- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt +++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt @@ -519,9 +519,9 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/ci/official/requirements_updater/nvidia-requirements.txt b/ci/official/requirements_updater/nvidia-requirements.txt index b979a97190a576..160e5f190a590b 100644 --- a/ci/official/requirements_updater/nvidia-requirements.txt +++ b/ci/official/requirements_updater/nvidia-requirements.txt @@ -12,11 +12,11 @@ nvidia-cufft-cu12>=11.2.3.61,<12.0 nvidia-curand-cu12>=10.3.6.82,<11.0 nvidia-cusolver-cu12>=11.6.3.83,<12.0 nvidia-cusparse-cu12>=12.5.1.3,<13.0 -nvidia-nccl-cu12>=2.29.7,<3.0 +nvidia-nccl-cu12>=2.30.7,<3.0 nvidia-nvjitlink-cu12>=12.5.82,<13.0 nvidia-nvshmem-cu12>=3.2.5 # CUDA 13 wheels -nvidia-nccl-cu13>=2.29.7 +nvidia-nccl-cu13>=2.30.7 nvidia-nvshmem-cu13>=3.3.20 nvidia-cublas>=13.0.0.19 nvidia-cuda-cupti>=13.0.48 diff --git a/ci/official/utilities/code_check_full.bats b/ci/official/utilities/code_check_full.bats index 1f6e6f810f91f7..b01432c7de0949 100644 --- a/ci/official/utilities/code_check_full.bats +++ b/ci/official/utilities/code_check_full.bats @@ -223,7 +223,7 @@ EOF --repo_env=HERMETIC_CUDA_VERSION="12.5.1" \ --repo_env=HERMETIC_CUDNN_VERSION="9.10.0" \ --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" \ - --repo_env=HERMETIC_NCCL_VERSION="2.29.7" \ + --repo_env=HERMETIC_NCCL_VERSION="2.30.7" \ "somepath(//tensorflow/tools/pip_package:wheel, " \ "@local_config_cuda//cuda:cudart + "\ "@local_config_cuda//cuda:cudart + "\ @@ -249,7 +249,7 @@ EOF --repo_env=HERMETIC_CUDA_VERSION="12.5.1" \ --repo_env=HERMETIC_CUDNN_VERSION="9.10.0" \ --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" \ - --repo_env=HERMETIC_NCCL_VERSION="2.29.7" \ + --repo_env=HERMETIC_NCCL_VERSION="2.30.7" \ --define framework_shared_object=false \ "somepath(//tensorflow/tools/pip_package:wheel, " \ "@local_config_cuda//cuda:cudart + "\ diff --git a/requirements_lock_3_10.txt b/requirements_lock_3_10.txt index 11f87a3b5687e2..e2298fa546a5bb 100644 --- a/requirements_lock_3_10.txt +++ b/requirements_lock_3_10.txt @@ -610,18 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_11.txt b/requirements_lock_3_11.txt index f1f748323495a2..df484dd450e71a 100644 --- a/requirements_lock_3_11.txt +++ b/requirements_lock_3_11.txt @@ -610,15 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_12.txt b/requirements_lock_3_12.txt index a412ddc80dbdb9..af658e13cac1e0 100644 --- a/requirements_lock_3_12.txt +++ b/requirements_lock_3_12.txt @@ -610,17 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_13.txt b/requirements_lock_3_13.txt index f860aee8c0739a..da604d3fa08573 100644 --- a/requirements_lock_3_13.txt +++ b/requirements_lock_3_13.txt @@ -610,17 +610,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/requirements_lock_3_14.txt b/requirements_lock_3_14.txt index 5bc65a2d12c992..0a4539230b6f6c 100644 --- a/requirements_lock_3_14.txt +++ b/requirements_lock_3_14.txt @@ -716,17 +716,16 @@ nvidia-cusparse-cu12==12.5.1.3 \ # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt # nvidia-cusolver-cu12 -nvidia-nccl-cu12==2.29.7 \ - --hash=sha256:0cf032ee22b560447daf0456108a75e32bd74a4de6c6b64725637a359fa48cd8 \ - --hash=sha256:ecd0a012051abc20c1aa87328841efa8cade3ced65803046e38c2f03c0891fea +nvidia-nccl-cu12==2.30.7 \ + --hash=sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1 \ + --hash=sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-requirements.txt -nvidia-nccl-cu13==2.29.7 \ - --hash=sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d \ - --hash=sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643 \ - --hash=sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 +nvidia-nccl-cu13==2.30.7 \ + --hash=sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77 \ + --hash=sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50 # via # -c ci/official/requirements_updater/nvidia-constraints.txt # -r ci/official/requirements_updater/nvidia-constraints.txt diff --git a/third_party/nccl/nccl_wheel.patch b/third_party/nccl/nccl_wheel.patch index f87b59ab0da5c2..30caf94090d360 100644 --- a/third_party/nccl/nccl_wheel.patch +++ b/third_party/nccl/nccl_wheel.patch @@ -1,48 +1,54 @@ diff --git a/include/nccl_device/gin/proxy/gin_proxy.h b/include/nccl_device/gin/proxy/gin_proxy.h +index 8a6d7dd2..78bd809e 100644 --- a/include/nccl_device/gin/proxy/gin_proxy.h +++ b/include/nccl_device/gin/proxy/gin_proxy.h -@@ -55,7 +55,7 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG --// 4x16 byte store with the write-through cache hint -+// 16x4 byte store with the write-through cache hint - #pragma unroll -- for (uint8_t i = 0; i < 4; i++) { -- __stwt((uint4*)&q[idx] + i, ((uint4*)gfd)[i]); +@@ -85,8 +85,8 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG + // ncclGinProxyGfd_t is declared __attribute__((packed, aligned(16))) in + // gin_proxy_device_host_common.h; static_asserts there enforce the contract. + NVCC_PRAGMA_UNROLL_AUTO +- for (uint8_t i = 0; i < sizeof(ncclGinProxyGfd_t) / sizeof(uint4); i++) { +- __stwt((uint4*)&q[gfdIdx] + i, ((uint4*)gfd)[i]); + for (uint8_t i = 0; i < 16; i++) { -+ __stwt((__half2*)&q[idx] + i, ((__half2*)gfd)[i]); ++ __stwt((__half2*)&q[gfdIdx] + i, ((__half2*)gfd)[i]); } - } - } + if (isGet) { + // Atomic max with rolling logic. diff --git a/include/nccl_device/utility.h b/include/nccl_device/utility.h +index 6e09fc10..dbc1cc28 100644 --- a/include/nccl_device/utility.h +++ b/include/nccl_device/utility.h @@ -17,7 +17,7 @@ - #define NCCL_CHECK_CUDACC 0 - #endif - #else -- #if __CUDACC__ -+ #ifdef __CUDACC__ - #define NCCL_CHECK_CUDACC 1 - #else - #define NCCL_CHECK_CUDACC 0 - #endif -diff --git a/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h b/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h ---- a/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h -+++ b/include/nccl_device/gin/gdaki/doca_gpunetio/common/doca_gpunetio_verbs_def.h -@@ -49,7 +49,7 @@ extern "C" { - /** - * Macro to temporarily cast a variable to volatile. - */ --#define DOCA_GPUNETIO_VOLATILE(x) (*(volatile typeof(x) *)&(x)) -+#define DOCA_GPUNETIO_VOLATILE(x) (*(volatile __typeof__(x) *)&(x)) - - /** - * Default warp size value of 32 threads -@@ -123,7 +123,7 @@ extern "C" { - (1ULL << DOCA_GPUNETIO_VERBS_MAX_TRANSFER_SIZE_SHIFT) // 1GiB + #define NCCL_CHECK_CUDACC 0 + #endif + #else +-#if __CUDACC__ ++#ifdef __CUDACC__ + #define NCCL_CHECK_CUDACC 1 + #else + #define NCCL_CHECK_CUDACC 0 +diff --git a/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h b/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h +index 4a2cf812..e38f7fc2 100644 +--- a/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h ++++ b/include/nccl_device/gin/gpi/gin_gpi_device_host_common.h +@@ -170,7 +170,7 @@ enum gpi_post_mode { + }; - #ifndef ACCESS_ONCE --#define ACCESS_ONCE(x) (*(volatile typeof(x) *)&(x)) -+#define ACCESS_ONCE(x) (*(volatile __typeof__(x) *)&(x)) + #ifndef GPI_ACCESS_ONCE +-#define GPI_ACCESS_ONCE(x) (*(volatile typeof(x)*)&(x)) ++#define GPI_ACCESS_ONCE(x) (*(volatile __typeof__(x)*)&(x)) #endif - #ifndef READ_ONCE \ No newline at end of file + #ifndef GPI_READ_ONCE +diff --git a/include/nccl_device/gin/gpi/gin_gpi.h b/include/nccl_device/gin/gpi/gin_gpi.h +index c004b4cb..7357c4c0 100644 +--- a/include/nccl_device/gin/gpi/gin_gpi.h ++++ b/include/nccl_device/gin/gpi/gin_gpi.h +@@ -212,7 +212,7 @@ __device__ static inline void gpi_gpu_channel_post_gfd_thread(gpi_gpu_channel_t* + // Manual PTX for MMIO 128-bit store (b128 needs CUDA 12.3+ / PTX 8.3) + uint64_t val_lo = segment[0].raw; + uint64_t val_hi = segment[1].raw; +-#if CUDART_VERSION >= 12030 ++#if CUDART_VERSION >= 12030 && !defined(__clang__) + asm volatile(R"YYY( + .reg .b128 _v%=; + mov.b128 _v%=, {%1, %2}; \ No newline at end of file diff --git a/third_party/nccl/workspace.bzl b/third_party/nccl/workspace.bzl index 38acb8b1693ce9..fb54ca1043572e 100644 --- a/third_party/nccl/workspace.bzl +++ b/third_party/nccl/workspace.bzl @@ -22,7 +22,7 @@ def repo(): name = "nccl_archive", build_file = "@xla//third_party/nccl:archive.BUILD", patch_file = ["@xla//third_party/nccl:archive.patch"], - sha256 = "e67239212c395bfdb398a7519491840d06fdf6b599c299f97c7ed0109777bba1", - strip_prefix = "nccl-2.29.7-1", - urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.29.7-1.tar.gz"), + sha256 = "292a7f7a27b6754acaf46b5506a60758ca7b18cc1dfbd3d1d4e1e229d0863b4e", + strip_prefix = "nccl-2.30.7-1", + urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.30.7-1.tar.gz"), ) diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc index 054ade6bdf69c5..0c18ed64c47782 100644 --- a/third_party/xla/tensorflow.bazelrc +++ b/third_party/xla/tensorflow.bazelrc @@ -189,7 +189,7 @@ common:mkl_aarch64 --config=mkl_aarch64_threadpool common:cuda_version --repo_env=HERMETIC_CUDA_VERSION="13.2.0" common:cuda_version --repo_env=HERMETIC_CUDNN_VERSION="9.12.0" common:cuda_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.3.20" -common:cuda_version --repo_env=HERMETIC_NCCL_VERSION="2.29.7" +common:cuda_version --repo_env=HERMETIC_NCCL_VERSION="2.30.7" # CUDA: This config refers to building CUDA op kernels with nvcc. common:cuda --repo_env TF_NEED_CUDA=1 diff --git a/third_party/xla/third_party/nccl/archive.BUILD b/third_party/xla/third_party/nccl/archive.BUILD index 5c8c2d0fd09687..70bb40f1017a2c 100644 --- a/third_party/xla/third_party/nccl/archive.BUILD +++ b/third_party/xla/third_party/nccl/archive.BUILD @@ -38,7 +38,7 @@ exports_files(["LICENSE.txt"]) NCCL_MAJOR = 2 -NCCL_MINOR = 29 +NCCL_MINOR = 30 NCCL_PATCH = 7 @@ -232,7 +232,6 @@ cc_library( # from the virtual includes directory. "src/include/collectives.h", "src/nccl.h", - "src/ras/ras_internal.h", ], hdrs = ["src/nccl.h"], include_prefix = "third_party/nccl", diff --git a/third_party/xla/third_party/nccl/archive.patch b/third_party/xla/third_party/nccl/archive.patch index 4b56bb6ca368d6..5844ea08fb9cab 100644 --- a/third_party/xla/third_party/nccl/archive.patch +++ b/third_party/xla/third_party/nccl/archive.patch @@ -2,29 +2,30 @@ diff --git a/src/device/common.cu b/src/device/common.cu.cc similarity index 100% rename from src/device/common.cu rename to src/device/common.cu.cc -diff --git a/src/device/onerank.cu b/src/device/onerank.cu.cc -similarity index 100% -rename from src/device/onerank.cu -rename to src/device/onerank.cu.cc diff --git a/src/device/common.h b/src/device/common.h +index 1c1f5d6f..700c27a7 100644 --- a/src/device/common.h +++ b/src/device/common.h -@@ -24,7 +24,7 @@ - #endif - - typedef void(*ncclDevFuncPtr_t)(); +@@ -26,9 +26,9 @@ + typedef void (*ncclDevFuncPtr_t)(); + #if defined(NCCL_OS_WINDOWS) + /* MSVC C2133: extern array of unknown size needs a complete type; use pointer instead. */ +-extern __device__ ncclDevFuncPtr_t const* ncclDevFuncTable; ++extern __device__ ncclDevFuncPtr_t* ncclDevFuncTable; + #else -extern __device__ ncclDevFuncPtr_t const ncclDevFuncTable[]; +extern __device__ ncclDevFuncPtr_t ncclDevFuncTable[]; - + #endif + struct ncclShmemGroup { - ncclConnInfo *recvConns[NCCL_MAX_ARITY]; diff --git a/src/device/generate.py b/src/device/generate.py +index 88ddcef9..73b88d2b 100755 --- a/src/device/generate.py +++ b/src/device/generate.py -@@ -209,8 +209,8 @@ kernel_funcs = sorted(set(best_kernel(*fn) for fn in primary_funcs)) - +@@ -207,8 +207,8 @@ kernel_funcs = sorted(set(best_kernel(*fn) for fn in primary_funcs)) + ################################################################################ - + -# Generate /device_table.cu -with open(os.path.join(gensrc, "device_table.cu"), "w") as f: +# Generate /device_table.cu.cc @@ -32,16 +33,19 @@ diff --git a/src/device/generate.py b/src/device/generate.py out = f.write out('#include "common.h"\n') out("\n") -@@ -225,7 +225,7 @@ with open(os.path.join(gensrc, "device_table.cu"), "w") as f: - out("#endif\n") - out("\n") - -- out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable[] = {\n"); -+ out("__device__ ncclDevFuncPtr_t ncclDevFuncTable[] = {\n"); +@@ -227,9 +227,9 @@ with open(os.path.join(gensrc, "device_table.cu"), "w") as f: + # an internal array + pointer alias. On Linux/GCC the array is named directly + # to avoid an extra device-memory indirection on every kernel dispatch. + out("#if defined(NCCL_OS_WINDOWS)\n") +- out("__device__ ncclDevFuncPtr_t const ncclDevFuncTableData[] = {\n") ++ out("__device__ ncclDevFuncPtr_t ncclDevFuncTableData[] = {\n") + out("#else\n") +- out("__device__ ncclDevFuncPtr_t const ncclDevFuncTable[] = {\n") ++ out("__device__ ncclDevFuncPtr_t ncclDevFuncTable[] = {\n") + out("#endif\n") index = 0 for fn in primary_funcs: - sym = paste("_", "ncclDevFunc", *fn) -@@ -283,8 +283,16 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: +@@ -292,8 +292,16 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: cudart, _ = required_cuda(*kfn) sym = paste("_", "ncclDevKernel", *kfn) if cudart != 0: out("#if CUDART_VERSION >= %d\n" % cudart) @@ -60,9 +64,9 @@ diff --git a/src/device/generate.py b/src/device/generate.py index += 1 out("nullptr};\n") out("\n") -@@ -298,15 +306,23 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: +@@ -307,15 +315,23 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: out("\n") - + # Maps primary id to kernel function pointer. - out("extern void* const ncclDevKernelForFunc[] = {\n") index = 0 @@ -87,16 +91,16 @@ diff --git a/src/device/generate.py b/src/device/generate.py index += 1 out("nullptr};\n") out("\n") -@@ -325,7 +341,7 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: +@@ -334,7 +350,7 @@ with open(os.path.join(gensrc, "host_table.cc"), "w") as f: # "coll" is reflected in the name: formally that no two funcs having different # coll's map to the same filename. def impl_filename(coll, redop, ty, algo, proto): - return "%s.cu" % paste("_", coll_camel_to_lower[coll], redop and redop.lower(), ty) + return "%s.cu.cc" % paste("_", coll_camel_to_lower[coll], redop and redop.lower(), ty) - + # Partition the functions and kernels to the .cu filenames. The partition is # a dictionary mapping filename to (coll, func-tuple list) -@@ -357,7 +373,7 @@ if os.environ.get("NCCL_USE_CMAKE", "0") != "1": +@@ -366,7 +382,7 @@ if os.environ.get("NCCL_USE_CMAKE", "0") != "1": with open(os.path.join(gensrc, "rules.mk"), "w") as f: out = f.write impl_names = sorted(name_to_funcs.keys()) @@ -105,66 +109,71 @@ diff --git a/src/device/generate.py b/src/device/generate.py out("LIB_OBJS_GEN = $(patsubst %,$(OBJDIR)/genobj/%.o,{names})\n" .format(names=" ".join(names))) out("\n") - +diff --git a/src/device/onerank.cu b/src/device/onerank.cu.cc +similarity index 100% +rename from src/device/onerank.cu +rename to src/device/onerank.cu.cc +diff --git a/src/device/symmetric/data_ops.cuh b/src/device/symmetric/data_ops.cuh +index 25b6a922..315b04d4 100644 +--- a/src/device/symmetric/data_ops.cuh ++++ b/src/device/symmetric/data_ops.cuh +@@ -41,7 +41,7 @@ static __device__ __forceinline__ T loadMem(GMemTag, T* p) { + for (int i = 0; i < sizeof(T) / 4; i++) u32[i] = __ldcs((uint32_t*)p + i); + break; + case 8: +- for (int i = 0; i < sizeof(T) / 8; i++) u64[i] = __ldcs((uint64_t*)p + i); ++ for (int i = 0; i < sizeof(T) / 8; i++) u64[i] = __ldcs((unsigned long long*)p + i); + break; + case 16: + for (int i = 0; i < sizeof(T) / 16; i++) u32v4[i] = __ldcs((uint4*)p + i); +@@ -89,7 +89,7 @@ static __device__ __forceinline__ void storeMem(GMemTag, T* p, T val) { + for (int i = 0; i < sizeof(T) / 4; i++) __stcs((uint32_t*)p + i, u32[i]); + break; + case 8: +- for (int i = 0; i < sizeof(T) / 8; i++) __stcs((uint64_t*)p + i, u64[i]); ++ for (int i = 0; i < sizeof(T) / 8; i++) __stcs((unsigned long long*)p + i, u64[i]); + break; + case 16: + for (int i = 0; i < sizeof(T) / 16; i++) __stcs((uint4*)p + i, u32v4[i]); diff --git a/src/include/nccl_common.h b/src/include/nccl_common.h +index 70b669a5..2667e055 100644 --- a/src/include/nccl_common.h +++ b/src/include/nccl_common.h -@@ -9,7 +9,9 @@ - #define NCCL_DEBUG_H_ - - // Workaround for libstdc++ trying to force public visibility of std:: symbols. We don't want to do that in libnccl.so. +@@ -11,7 +11,9 @@ + #ifdef NCCL_OS_LINUX + // Workaround for libstdc++ trying to force public visibility of std:: symbols. We don't want to do that in + // libnccl.so. +#if defined(__GLIBCXX__) #include +#endif #undef _GLIBCXX_VISIBILITY #define _GLIBCXX_VISIBILITY(V) - + #endif diff --git a/src/include/nccl_device/gin/proxy/gin_proxy.h b/src/include/nccl_device/gin/proxy/gin_proxy.h +index 8a6d7dd2..78bd809e 100644 --- a/src/include/nccl_device/gin/proxy/gin_proxy.h +++ b/src/include/nccl_device/gin/proxy/gin_proxy.h -@@ -55,7 +55,7 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG --// 4x16 byte store with the write-through cache hint -+// 16x4 byte store with the write-through cache hint - #pragma unroll -- for (uint8_t i = 0; i < 4; i++) { -- __stwt((uint4*)&q[idx] + i, ((uint4*)gfd)[i]); +@@ -85,8 +85,8 @@ NCCL_DEVICE_INLINE void postGfd(Coop coop, ncclGinProxyGpuCtx_t* proxyCtx, ncclG + // ncclGinProxyGfd_t is declared __attribute__((packed, aligned(16))) in + // gin_proxy_device_host_common.h; static_asserts there enforce the contract. + NVCC_PRAGMA_UNROLL_AUTO +- for (uint8_t i = 0; i < sizeof(ncclGinProxyGfd_t) / sizeof(uint4); i++) { +- __stwt((uint4*)&q[gfdIdx] + i, ((uint4*)gfd)[i]); + for (uint8_t i = 0; i < 16; i++) { -+ __stwt((__half2*)&q[idx] + i, ((__half2*)gfd)[i]); ++ __stwt((__half2*)&q[gfdIdx] + i, ((__half2*)gfd)[i]); } - } - } -diff --git a/src/include/nccl_device/utility.h b/src/include/nccl_device/utility.h ---- a/src/include/nccl_device/utility.h -+++ b/src/include/nccl_device/utility.h -@@ -17,7 +17,7 @@ - #define NCCL_CHECK_CUDACC 0 - #endif - #else -- #if __CUDACC__ -+ #ifdef __CUDACC__ - #define NCCL_CHECK_CUDACC 1 - #else - #define NCCL_CHECK_CUDACC 0 - #endif - -diff --git a/src/device/symmetric/data_ops.cuh b/src/device/symmetric/data_ops.cuh ---- a/src/device/symmetric/data_ops.cuh -+++ b/src/device/symmetric/data_ops.cuh -@@ -26,7 +26,7 @@ - case 1: for (int i=0; i < sizeof(T)/1; i++) u8[i] = __ldcs((uint8_t*)p + i); break; - case 2: for (int i=0; i < sizeof(T)/2; i++) u16[i] = __ldcs((uint16_t*)p + i); break; - case 4: for (int i=0; i < sizeof(T)/4; i++) u32[i] = __ldcs((uint32_t*)p + i); break; -- case 8: for (int i=0; i < sizeof(T)/8; i++) u64[i] = __ldcs((uint64_t*)p + i); break; -+ case 8: for (int i=0; i < sizeof(T)/8; i++) u64[i] = __ldcs((unsigned long long*)p + i); break; - case 16: for (int i=0; i < sizeof(T)/16; i++) u32v4[i] = __ldcs((uint4*)p + i); break; - default: __builtin_unreachable(); - } -@@ -62,7 +62,7 @@ - case 1: for (int i=0; i < sizeof(T)/1; i++) __stcs((uint8_t*)p + i, u8[i]); break; - case 2: for (int i=0; i < sizeof(T)/2; i++) __stcs((uint16_t*)p + i, u16[i]); break; - case 4: for (int i=0; i < sizeof(T)/4; i++) __stcs((uint32_t*)p + i, u32[i]); break; -- case 8: for (int i=0; i < sizeof(T)/8; i++) __stcs((uint64_t*)p + i, u64[i]); break; -+ case 8: for (int i=0; i < sizeof(T)/8; i++) __stcs((unsigned long long*)p + i, u64[i]); break; - case 16: for (int i=0; i < sizeof(T)/16; i++) __stcs((uint4*)p + i, u32v4[i]); break; - default: __builtin_unreachable(); - } + if (isGet) { + // Atomic max with rolling logic. +diff --git a/src/include/nccl_device/gin/gpi/gin_gpi.h b/src/include/nccl_device/gin/gpi/gin_gpi.h +index c004b4cb..7357c4c0 100644 +--- a/src/include/nccl_device/gin/gpi/gin_gpi.h ++++ b/src/include/nccl_device/gin/gpi/gin_gpi.h +@@ -212,7 +212,7 @@ __device__ static inline void gpi_gpu_channel_post_gfd_thread(gpi_gpu_channel_t* + // Manual PTX for MMIO 128-bit store (b128 needs CUDA 12.3+ / PTX 8.3) + uint64_t val_lo = segment[0].raw; + uint64_t val_hi = segment[1].raw; +-#if CUDART_VERSION >= 12030 ++#if CUDART_VERSION >= 12030 && !defined(__clang__) + asm volatile(R"YYY( + .reg .b128 _v%=; + mov.b128 _v%=, {%1, %2}; \ No newline at end of file diff --git a/third_party/xla/third_party/nccl/workspace.bzl b/third_party/xla/third_party/nccl/workspace.bzl index dd40c3264966d7..ba05442b846048 100644 --- a/third_party/xla/third_party/nccl/workspace.bzl +++ b/third_party/xla/third_party/nccl/workspace.bzl @@ -22,7 +22,7 @@ def repo(): name = "nccl_archive", build_file = "//third_party/nccl:archive.BUILD", patch_file = ["//third_party/nccl:archive.patch"], - sha256 = "e67239212c395bfdb398a7519491840d06fdf6b599c299f97c7ed0109777bba1", - strip_prefix = "nccl-2.29.7-1", - urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.29.7-1.tar.gz"), + sha256 = "292a7f7a27b6754acaf46b5506a60758ca7b18cc1dfbd3d1d4e1e229d0863b4e", + strip_prefix = "nccl-2.30.7-1", + urls = tf_mirror_urls("https://github.com/NVIDIA/nccl/archive/refs/tags/v2.30.7-1.tar.gz"), ) From 722e8e16fafa6c4dd36cc33496591b8523c6ba87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= Date: Fri, 28 Aug 2026 02:49:30 -0700 Subject: [PATCH 17/23] [Mosaic] Add attribute to identify main function PiperOrigin-RevId: 972463068 --- third_party/xla/xla/mosaic/dialect/tpu/tpu.td | 4 ++- .../xla/xla/mosaic/dialect/tpu/tpu_dialect.cc | 36 +++++++++++++++++-- .../xla/xla/mosaic/dialect/tpu/tpu_dialect.h | 5 ++- 3 files changed, 39 insertions(+), 6 deletions(-) diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td index 0231285f239bc2..d4e241c0f16f3e 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu.td @@ -25,12 +25,14 @@ def TPU_Dialect : Dialect { let useDefaultAttributePrinterParser = 1; let useDefaultTypePrinterParser = 1; let extraClassDeclaration = [{ - static StringRef GetCoreTypeKey() { return "tpu.core_type"; } + static constexpr StringRef GetMainKey() { return "tpu.main"; } + static constexpr StringRef GetCoreTypeKey() { return "tpu.core_type"; } static std::optional GetCoreTypeAttr(Operation *op); }]; let hasConstantMaterializer = 1; let hasCanonicalizer = 1; + let hasOperationAttrVerify = 1; } class TPU_Attr traits = []> diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc index 0a9e036b37033e..3064f345c4710b 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc @@ -43,6 +43,7 @@ limitations under the License. #include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/AffineExpr.h" #include "mlir/IR/AffineMap.h" +#include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypeInterfaces.h" @@ -115,6 +116,19 @@ Operation* TPUDialect::materializeConstant(OpBuilder& builder, Attribute value, return mlir::cast(attr).getValue(); } +LogicalResult TPUDialect::verifyOperationAttribute(Operation* const op, + NamedAttribute attr) { + CHECK(op != nullptr); + if (attr.getName() == GetMainKey()) { + if (!isa(op) || !isa(attr.getValue())) { + return op->emitOpError() << GetMainKey() + << " attribute is expected to be a unit " + "attribute on a func.func operation"; + } + } + return success(); +} + // Rewrites // // memref.dim(tpu.memref_slice(..., dynamic_sizes), i) @@ -216,9 +230,27 @@ CoreType GetCoreTypeOfParentOp(Operation& op) { return parent ? *TPUDialect::GetCoreTypeAttr(parent) : CoreType::kTc; } -absl::StatusOr GetFuncWithCoreType(ModuleOp module, - CoreType core_type) { +absl::StatusOr GetMainFunc(ModuleOp module, CoreType core_type) { func::FuncOp result = nullptr; + for (auto func_op : module.getOps()) { + if (!func_op->hasAttr(TPUDialect::GetMainKey())) { + continue; + } + if (TPUDialect::GetCoreTypeAttr(func_op) != core_type) { + continue; + } + if (result != nullptr) { + return absl::InvalidArgumentError(absl::StrFormat( + "Multiple functions with %v attribute and tpu.core_type = %v found", + TPUDialect::GetMainKey(), core_type)); + } + result = func_op; + } + if (result != nullptr) { + return result; + } + // If there is no function marked with tpu.main, fall back to finding a + // (unique) function with the requested core type. for (func::FuncOp func_op : module.getOps()) { if (TPUDialect::GetCoreTypeAttr(func_op) != core_type) { continue; diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h index a488ee1d7c98b5..3f886e4a1a2347 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.h @@ -84,9 +84,8 @@ Operation* GetParentOpWithCoreType(Operation& op); // annotation is found, returns kTc. CoreType GetCoreTypeOfParentOp(Operation& op); -// Returns the function in the module with the given core type. -absl::StatusOr GetFuncWithCoreType(ModuleOp module, - CoreType core_type); +// Returns the main function in the module for the given core type. +absl::StatusOr GetMainFunc(ModuleOp module, CoreType core_type); // Changes the memory space of the value and propagates it through the program. LogicalResult specializeMemorySpace(TypedValue value, From 8ab2e8fe5d534ecbdb7673bd99537974e757f8db Mon Sep 17 00:00:00 2001 From: Ilya Tikhonovskiy Date: Fri, 28 Aug 2026 02:53:44 -0700 Subject: [PATCH 18/23] Allow mixed FP8 operand types in ConvolutionOp and DynamicConvOp. This enables lowering of mixed FP8 dot operations to convolution. PiperOrigin-RevId: 972464941 --- .../xla/third_party/stablehlo/temporary.patch | 340 ++++++++++++++++++ .../hlo_to_mhlo/hlo_function_importer.cc | 10 +- 2 files changed, 347 insertions(+), 3 deletions(-) diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index fc07bbc2c695a5..dc2c4735acf69c 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -1,3 +1,41 @@ +diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md +--- stablehlo/docs/spec.md ++++ stablehlo/docs/spec.md +@@ -2493,7 +2493,9 @@ + * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. + * (C26) `rank(result) = N`. + * If the operation uses non-quantized tensors: +- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. ++ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or ++ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and ++ `element_type(result) = element_type(lhs)`). + * If the operation uses quantized tensors: + * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. + * (C29) If `is_per_axis_quantized(rhs)`, +@@ -3163,7 +3165,9 @@ + * `num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1`. + * (C26) `rank(result) = N`. + * If the operation uses non-quantized tensors: +- * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)`. ++ * (C27) `element_type(lhs) = element_type(rhs) = element_type(result)` or ++ (`is_fp8(element_type(lhs))` and `is_fp8(element_type(rhs))` and ++ `element_type(result) = element_type(lhs)`). + * If the operation uses quantized tensors: + * (C28) `is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)`. + * (C29) If `is_per_axis_quantized(rhs)`, +@@ -7654,6 +7658,12 @@ + * `is_quantized(x: Value | Placeholder | Type) -> Value` is a shortcut for + `is_quantized_tensor_element_type(x)`. + ++* `is_fp8(x: Value | Placeholder | Type) -> Value` returns `true` if `x` is one ++of `Float8E4M3Type`, `Float8E4M3FNType`, `Float8E4M3B11FNUZType`, ++`Float8E4M3FNUZType`, `Float8E5M2Type`, `Float8E5M2FNUZType`, `Float8E3M4Type`, ++or `Float8E8M0FNUType`. If `x` is a value or placeholder, this function is a ++shortcut for `is_fp8(type(x))`. ++ + * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all + types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`. + If `x` is a value or placeholder, this function is a shortcut for diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp --- stablehlo/stablehlo/dialect/Serialization.cpp +++ stablehlo/stablehlo/dialect/Serialization.cpp @@ -67,6 +105,144 @@ diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo return nullptr; } +diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp +--- stablehlo/stablehlo/dialect/TypeInference.cpp ++++ stablehlo/stablehlo/dialect/TypeInference.cpp +@@ -86,6 +86,13 @@ + // Utils for quantization specific verifications + //===----------------------------------------------------------------------===// + ++bool isFp8Type(mlir::Type type) { ++ return llvm::isa(type); ++} ++ + template + bool allQuantized(ArrayRef typeRange) { + return llvm::all_of( +@@ -2221,11 +2228,14 @@ + // convolution_c27 + if (!anyQuantized({rankedLhsType, rankedRhsType}) && + !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), +- rankedRhsType.getElementType())) ++ rankedRhsType.getElementType()) && ++ !(isFp8Type(rankedLhsType.getElementType()) && ++ isFp8Type(rankedRhsType.getElementType()))) { + return emitOptionalError( + location, "expects lhs and rhs to have compatible element type. Got: ", + rankedLhsType.getElementType(), " and ", + rankedRhsType.getElementType()); ++ } + + if (failed(verifyConvolutionAttributes( + location, lhsType, rhsType, inputBatchDimension, +@@ -2548,7 +2558,9 @@ + // dynamic_conv_c27 + if (!anyQuantized({rankedLhsType, rankedRhsType}) && + !isCompatibleForHloTypeInference(rankedLhsType.getElementType(), +- rankedRhsType.getElementType())) ++ rankedRhsType.getElementType()) && ++ !(isFp8Type(rankedLhsType.getElementType()) && ++ isFp8Type(rankedRhsType.getElementType()))) + return emitOptionalError( + location, "expects lhs and rhs to have compatible element type. Got: ", + rankedLhsType.getElementType(), " and ", +diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect/Version.h +--- stablehlo/stablehlo/dialect/Version.h ++++ stablehlo/stablehlo/dialect/Version.h +@@ -38,7 +38,7 @@ + static FailureOr fromString(llvm::StringRef versionRef); + + /// Return a Version representing the current VHLO dialect version. +- static Version getCurrentVersion() { return Version(1, 19, 0); } ++ static Version getCurrentVersion() { return Version(1, 20, 0); } + + /// Return a Version representing the minimum supported VHLO dialect version. + static Version getMinimumVersion() { return Version(0, 9, 0); } +diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td +--- stablehlo/stablehlo/dialect/VhloDialect.td ++++ stablehlo/stablehlo/dialect/VhloDialect.td +@@ -58,6 +58,7 @@ + 1.17.0: Add `future` type support to `custom_call` op. + 1.18.0: Add `result_tilings` attribute to `custom_call` op. + 1.19.0: Add CollectiveReduceOp. ++ 1.20.0: Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops. + }]; + + let useDefaultAttributePrinterParser = 0; +diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.cpp b/stablehlo/stablehlo/dialect/VhloOps.cpp +--- stablehlo/stablehlo/dialect/VhloOps.cpp ++++ stablehlo/stablehlo/dialect/VhloOps.cpp +@@ -359,11 +359,44 @@ + return success(); + } + ++bool isVhloFp8Type(Type type) { ++ return isa(type); ++} ++ ++LogicalResult verifyConstraint_1_20_0(mlir::Operation* op, ++ Version targetVersion) { ++ if (targetVersion < Version(1, 20, 0)) { ++ if (op->getNumOperands() < 2) { ++ return failure(); ++ } ++ Type lhsElementType = getVhloElementType(op->getOperand(0).getType()); ++ Type rhsElementType = getVhloElementType(op->getOperand(1).getType()); ++ if (lhsElementType != rhsElementType && isVhloFp8Type(lhsElementType) && ++ isVhloFp8Type(rhsElementType)) { ++ return failure(); ++ } ++ } ++ return success(); ++} ++ + } // namespace + + LogicalResult AllReduceOpV1::validateConstraint(mlir::Operation* op, + Version targetVersion) { + return verifyConstraint_0_17_0(op, targetVersion); ++} ++ ++LogicalResult ConvolutionOpV1::validateConstraint(mlir::Operation* op, ++ Version targetVersion) { ++ return verifyConstraint_1_20_0(op, targetVersion); ++} ++ ++LogicalResult DynamicConvOpV2::validateConstraint(mlir::Operation* op, ++ Version targetVersion) { ++ return verifyConstraint_1_20_0(op, targetVersion); + } + + LogicalResult ReduceOpV1::validateConstraint(mlir::Operation* op, +diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.td b/stablehlo/stablehlo/dialect/VhloOps.td +--- stablehlo/stablehlo/dialect/VhloOps.td ++++ stablehlo/stablehlo/dialect/VhloOps.td +@@ -390,7 +390,8 @@ + let results = (outs VHLO_AnyType:$result); + } + +-def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current"> { ++def VHLO_ConvolutionOpV1 : VHLO_Op<"convolution_v1", "0.9.0", "current", ++ [DeclareOpInterfaceMethods]> { + let arguments = (ins + VHLO_AnyType:$lhs, + VHLO_AnyType:$rhs, +@@ -572,7 +573,8 @@ + + // Padding should be specified as an operand only, not an attribute. + // Remove `d_padding` and convert `padding` to an operand. +-def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current"> { ++def VHLO_DynamicConvOpV2 : VHLO_Op<"dynamic_conv_v2", "0.20.0", "current", ++ [DeclareOpInterfaceMethods]> { + let arguments = (ins + VHLO_AnyType:$lhs, + VHLO_AnyType:$rhs, diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp --- stablehlo/stablehlo/tests/TestUtils.cpp +++ stablehlo/stablehlo/tests/TestUtils.cpp @@ -312,6 +488,170 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stableh + return %0 : tensor<3x4x5xf64> +} + +diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir +--- stablehlo/stablehlo/tests/verify_convolution.mlir ++++ stablehlo/stablehlo/tests/verify_convolution.mlir +@@ -78,6 +78,60 @@ + + // ----- + ++// CHECK-LABEL: func @convolution_mixed_fp8 ++func.func @convolution_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, ++ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { ++ %result = "stablehlo.convolution"(%arg0, %arg1) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ padding = dense<2> : tensor<2x2xi64>, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> ++ tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++} ++ ++// ----- ++ ++func.func @convolution_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { ++ // expected-error@+1{{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} ++ %result = "stablehlo.convolution"(%arg0, %arg1) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ padding = dense<2> : tensor<2x2xi64>, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> ++ func.return %result : tensor<100x28x28x1xf32> ++} ++ ++// ----- ++ + func.func @convolution(%arg0: tensor<2x2x3x4xf32>, %arg1: tensor<3x5x5x3xf32>) -> tensor<3x5x5x4xf32> { + // expected-error@+3{{Unexpected keyword stide}} + %0 = stablehlo.convolution(%arg0, %arg1) +diff --ruN a/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir b/stablehlo/stablehlo/tests/verify_dynamic_conv.mlir +--- stablehlo/stablehlo/tests/verify_dynamic_conv.mlir ++++ stablehlo/stablehlo/tests/verify_dynamic_conv.mlir +@@ -48,6 +48,36 @@ + + // ----- + ++// CHECK-LABEL: func @dynamic_conv_mixed_fp8 ++func.func @dynamic_conv_mixed_fp8(%arg0 : tensor<100x26x26x32xf8E5M2>, ++ %arg1 : tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { ++ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> ++ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { ++ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, ++ feature_group_count = 1 : i64, ++ batch_group_count = 1 : i64 ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> ++ tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++} ++ ++// ----- ++ ++func.func @dynamic_conv_mismatched_element_types(%arg0: tensor<100x26x26x32xf32>, ++ %arg1: tensor<3x3x1x32xf16>) -> tensor<100x28x28x1xf32> { ++ // expected-error@+2 {{expects lhs and rhs to have compatible element type. Got: 'f32' and 'f16'}} ++ %padding = stablehlo.constant dense<2> : tensor<2x2xi64> ++ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %padding) { ++ dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>, ++ feature_group_count = 1 : i64, ++ batch_group_count = 1 : i64 ++ } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf16>, tensor<2x2xi64>) -> ++ tensor<100x28x28x1xf32> ++ func.return %result : tensor<100x28x28x1xf32> ++} ++ ++// ----- ++ + func.func @dynamic_conv_c1(%arg0: tensor<1x8x8x207xf32>, + %arg1: tensor<3x3x207xf32>) -> tensor<1x8x8x16xf32> { + // expected-error@+2 {{expects convolution arguments to have same number of dimensions. Got: 'tensor<1x8x8x207xf32>' and 'tensor<3x3x207xf32>'.}} +diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir +--- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir ++++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_19_0.mlir +@@ -0,0 +1,56 @@ ++// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.19.0' --verify-diagnostics --split-input-file %s ++ ++// expected-error @+1 {{failed to convert VHLO to v1.19.0}} ++module { ++ func.func @convolution_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> { ++ // expected-error @+1 {{failed to legalize operation 'vhlo.convolution_v1' that was explicitly marked illegal}} ++ %result = "stablehlo.convolution"(%arg0, %arg1) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ padding = dense<2> : tensor<2x2xi64>, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>) -> tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++ } ++} ++ ++// ----- ++ ++// expected-error @+1 {{failed to convert VHLO to v1.19.0}} ++module { ++ func.func @dynamic_conv_mixed_fp8(%arg0: tensor<100x26x26x32xf8E5M2>, %arg1: tensor<3x3x1x32xf8E4M3FN>, %arg2: tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> { ++ // expected-error @+1 {{failed to legalize operation 'vhlo.dynamic_conv_v2' that was explicitly marked illegal}} ++ %result = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) { ++ batch_group_count = 1 : i64, ++ dimension_numbers = #stablehlo.conv, ++ feature_group_count = 1 : i64, ++ lhs_dilation = array, ++ rhs_dilation = array, ++ window_strides = array ++ } : (tensor<100x26x26x32xf8E5M2>, tensor<3x3x1x32xf8E4M3FN>, tensor<2x2xi64>) -> tensor<100x28x28x1xf8E5M2> ++ func.return %result : tensor<100x28x28x1xf8E5M2> ++ } ++} diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index e148e25c1065e4..cb64f9eb0defb3 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -2143,9 +2143,12 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( auto lhs_element_type = instruction->operand(0)->shape().element_type(); auto rhs_element_type = instruction->operand(1)->shape().element_type(); if (lhs_element_type != rhs_element_type) { - // Cast LHS or RHS to the common element type. - if (primitive_util::CastPreservesValues(lhs_element_type, - rhs_element_type)) { + if (primitive_util::IsF8Type(lhs_element_type) && + primitive_util::IsF8Type(rhs_element_type)) { + // Allow mixed FP8 without conversion. + } else if (primitive_util::CastPreservesValues(lhs_element_type, + rhs_element_type)) { + // Cast LHS to the common element type. auto convert_op_return_type = mlir::cast(lhs.getType()) .clone(mlir::getElementTypeOrSelf(rhs)); @@ -2153,6 +2156,7 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( convert_op_return_type, lhs); } else if (primitive_util::CastPreservesValues(rhs_element_type, lhs_element_type)) { + // Cast RHS to the common element type. auto convert_op_return_type = mlir::cast(rhs.getType()) .clone(mlir::getElementTypeOrSelf(lhs)); From f18f7c54419d8c64070ea8387d7e2de76a0f154b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 03:18:57 -0700 Subject: [PATCH 19/23] Rollback of PR #126177 Reverts d9a8da74b4c3de28a39ab34ad007838d6bc30c67 PiperOrigin-RevId: 972476125 --- tensorflow/lite/kernels/slice.cc | 11 +---- tensorflow/lite/kernels/slice_test.cc | 67 --------------------------- 2 files changed, 1 insertion(+), 77 deletions(-) diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc index 597c14f585d25c..86dca38bb7c427 100644 --- a/tensorflow/lite/kernels/slice.cc +++ b/tensorflow/lite/kernels/slice.cc @@ -145,16 +145,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size)); // If the shape of output is fully specified then resize even if // the input shape is not staticly defined. - // - // A fully specified output shape does not imply the slice is in bounds when - // the input extent is only known at run time. Taking this path leaves the - // output static, so Eval() -- which only re-runs ResizeOutputShape() for a - // dynamic output -- never reaches CalculateOutputShapeVector(), the one place - // `begin` and `size` are checked against the input. Fall through when the - // input has an unspecified dimension so the output is marked dynamic and the - // bounds are validated against the actual extent on every invocation. - if (ShapeHasRank(output->dims) && !HasUnspecifiedDimension(output) && - ShapeHasRank(input->dims) && !HasUnspecifiedDimension(input)) { + if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims)) { return kTfLiteOk; } // Postpone allocation of output if any of the indexing tensors is not diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc index e3b55f6116e32c..09b68642fff848 100644 --- a/tensorflow/lite/kernels/slice_test.cc +++ b/tensorflow/lite/kernels/slice_test.cc @@ -102,73 +102,6 @@ class SliceOpModel : public SingleOpModel { class SliceOpTest : public ::testing::TestWithParam {}; -// Model with a dynamic input dimension and a statically shaped output. The -// suite name is deliberately distinct from SliceOpTest, which is a TEST_P -// fixture -- gtest rejects a suite that mixes TEST and TEST_P. -class DynamicInputSliceOpModel : public SingleOpModel { - public: - DynamicInputSliceOpModel(TensorData input_data, - std::initializer_list begin_shape, - std::initializer_list begin_data, - std::initializer_list size_shape, - std::initializer_list size_data, - TensorData output_data) { - input_ = AddInput(input_data); - begin_ = AddConstInput(TensorType_INT32, begin_data, begin_shape); - size_ = AddConstInput(TensorType_INT32, size_data, size_shape); - output_ = AddOutput(output_data); - SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions, - CreateSliceOptions(builder_).Union()); - // Delegates are bypassed: a delegate that claims the SLICE node would set - // the output allocation type itself, so the assertions below would no - // longer describe the built-in CPU kernel. - BuildInterpreter({input_data.shape, begin_shape, size_shape}, - /*num_threads=*/-1, /*allow_fp32_relax_to_fp16=*/false, - /*apply_delegate=*/false); - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - const TfLiteTensor* GetOutputTensor() { - return interpreter_->tensor(output_); - } - - private: - int input_; - int begin_; - int size_; - int output_; -}; - -// A dynamic input dimension must force the output dynamic so the bounds are -// re-validated in Eval(), even though the declared output shape is static. -TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputValid) { - TensorData input_data(TensorType_FLOAT32, {1, 8}); - input_data.shape_signature = {1, -1}; - TensorData output_data(TensorType_FLOAT32, {1, 4}); - - DynamicInputSliceOpModel m(input_data, {2}, {0, 2}, {2}, {1, 4}, output_data); - m.SetInput({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); - ASSERT_EQ(m.Invoke(), kTfLiteOk); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4})); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.0, 4.0, 5.0, 6.0})); - EXPECT_EQ(m.GetOutputTensor()->allocation_type, kTfLiteDynamic); -} - -// The same path must reject a window that does not fit the actual extent. -TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputOutOfBounds) { - TensorData input_data(TensorType_FLOAT32, {1, 4}); - input_data.shape_signature = {1, -1}; - TensorData output_data(TensorType_FLOAT32, {1, 8}); - - DynamicInputSliceOpModel m(input_data, {2}, {0, 0}, {2}, {1, 8}, output_data); - m.SetInput({1.0, 2.0, 3.0, 4.0}); - EXPECT_EQ(m.Invoke(), kTfLiteError); -} - TEST_P(SliceOpTest, In1D) { SliceOpModel m({4}, {1}, {1}, {1}, {2}, TensorType_INT32, TensorType_FLOAT32, GetParam()); From 5db00615b5d71a554a6588946fd0bfd8a90b0069 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 03:21:32 -0700 Subject: [PATCH 20/23] Move `Singleprint(export_dir)` from `singleprint` to `fingerprinting`. PiperOrigin-RevId: 972477162 --- tensorflow/cc/saved_model/BUILD | 23 +++------------ tensorflow/cc/saved_model/fingerprinting.cc | 7 +++++ tensorflow/cc/saved_model/fingerprinting.h | 5 ++++ tensorflow/cc/saved_model/singleprint.cc | 31 --------------------- tensorflow/cc/saved_model/singleprint.h | 3 -- 5 files changed, 16 insertions(+), 53 deletions(-) diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD index f604a97fe95c8d..93b4bf515cbde8 100644 --- a/tensorflow/cc/saved_model/BUILD +++ b/tensorflow/cc/saved_model/BUILD @@ -518,18 +518,8 @@ cc_library( "//tensorflow/python:__pkg__", ], deps = [ - ":constants", "//tensorflow/core/protobuf:for_core_protos_cc", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@tsl//tsl/platform:errors", - "@tsl//tsl/platform:statusor", - ] + if_not_mobile([ - "//tensorflow/core:lib", - ]) + if_android([ - "//tensorflow/core:portable_tensorflow_lib_lite", - ]), + ], alwayslink = True, ) @@ -539,14 +529,9 @@ cc_library( visibility = ["//visibility:public"], deps = if_static([ ":singleprint_impl", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", + ]) + [ "//tensorflow/core/protobuf:for_core_protos_cc", - ]) + if_not_mobile([ - "//tensorflow/core:lib", - ]) + if_android([ - "//tensorflow/core:portable_tensorflow_lib_lite", - ]), + ], ) tf_cc_test( @@ -555,7 +540,7 @@ tf_cc_test( srcs = ["singleprint_test.cc"], deps = [ ":singleprint", - "//tensorflow/core:protos_all_cc", + "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_googletest//:gtest_main", ], ) diff --git a/tensorflow/cc/saved_model/fingerprinting.cc b/tensorflow/cc/saved_model/fingerprinting.cc index 9f7519130e6add..b2622a32cb0847 100644 --- a/tensorflow/cc/saved_model/fingerprinting.cc +++ b/tensorflow/cc/saved_model/fingerprinting.cc @@ -27,6 +27,7 @@ limitations under the License. #include "absl/strings/strip.h" #include "tensorflow/cc/saved_model/constants.h" #include "tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h" +#include "tensorflow/cc/saved_model/singleprint.h" #include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/graph/regularization/simple_delete.h" #include "tensorflow/core/graph/regularization/util.h" @@ -259,4 +260,10 @@ absl::StatusOr ReadSavedModelFingerprint( return fingerprint_proto; } +absl::StatusOr Singleprint(absl::string_view export_dir) { + TF_ASSIGN_OR_RETURN(FingerprintDef fingerprint_def, + ReadSavedModelFingerprint(export_dir)); + return Singleprint(fingerprint_def); +} + } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/cc/saved_model/fingerprinting.h b/tensorflow/cc/saved_model/fingerprinting.h index 25adbbe23799da..a4421874d3d2df 100644 --- a/tensorflow/cc/saved_model/fingerprinting.h +++ b/tensorflow/cc/saved_model/fingerprinting.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_ #define TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_ +#include + #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "tensorflow/cc/saved_model/singleprint.h" // IWYU pragma: keep. Used for Singleprint(). @@ -33,6 +35,9 @@ absl::StatusOr CreateFingerprintDef( absl::StatusOr ReadSavedModelFingerprint( absl::string_view export_dir); +// Canonical fingerprinting ID for a SavedModel loaded from `export_dir`. +absl::StatusOr Singleprint(absl::string_view export_dir); + } // namespace tensorflow::saved_model::fingerprinting #endif // TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_ diff --git a/tensorflow/cc/saved_model/singleprint.cc b/tensorflow/cc/saved_model/singleprint.cc index 4ffc3f533c90f5..a189c8a0293f45 100644 --- a/tensorflow/cc/saved_model/singleprint.cc +++ b/tensorflow/cc/saved_model/singleprint.cc @@ -18,35 +18,10 @@ limitations under the License. #include #include -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "tensorflow/cc/saved_model/constants.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/path.h" #include "tensorflow/core/protobuf/fingerprint.pb.h" -#include "tsl/platform/statusor.h" namespace tensorflow::saved_model::fingerprinting { -namespace { - -absl::StatusOr ReadSavedModelFingerprint( - absl::string_view export_dir) { - const std::string fingerprint_pb_path = - io::JoinPath(export_dir, kFingerprintFilenamePb); - TF_RETURN_IF_ERROR(Env::Default()->FileExists(fingerprint_pb_path)); - - FingerprintDef fingerprint_proto; - absl::Status result = - ReadBinaryProto(Env::Default(), fingerprint_pb_path, &fingerprint_proto); - if (!result.ok()) return result; - - return fingerprint_proto; -} - -} // namespace - std::string Singleprint(uint64_t graph_def_program_hash, uint64_t signature_def_hash, uint64_t saved_object_graph_hash, @@ -63,10 +38,4 @@ std::string Singleprint(const FingerprintDef& fingerprint) { fingerprint.saved_object_graph_hash(), fingerprint.checkpoint_hash()); } -absl::StatusOr Singleprint(absl::string_view export_dir) { - TF_ASSIGN_OR_RETURN(FingerprintDef fingerprint_def, - ReadSavedModelFingerprint(export_dir)); - return Singleprint(fingerprint_def); -} - } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/cc/saved_model/singleprint.h b/tensorflow/cc/saved_model/singleprint.h index 95e27d4f6de655..a83de66152ca3a 100644 --- a/tensorflow/cc/saved_model/singleprint.h +++ b/tensorflow/cc/saved_model/singleprint.h @@ -19,8 +19,6 @@ limitations under the License. #include #include -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" #include "tensorflow/core/protobuf/fingerprint.pb.h" namespace tensorflow::saved_model::fingerprinting { @@ -31,7 +29,6 @@ std::string Singleprint(uint64_t graph_def_program_hash, uint64_t saved_object_graph_hash, uint64_t checkpoint_hash); std::string Singleprint(const FingerprintDef& fingerprint); -absl::StatusOr Singleprint(absl::string_view export_dir); } // namespace tensorflow::saved_model::fingerprinting From 8fe8cca0c19e4f144f52301f1a8dad6af8965888 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 03:31:20 -0700 Subject: [PATCH 21/23] Adding `xla.memory_space_assignment.MsaTensorOverrides` to explicitly force instructions to be pinned, kept in default memory, or operands prefetched starting at a specified point. PiperOrigin-RevId: 972481193 --- .../xla/service/memory_space_assignment/BUILD | 5 +- .../memory_space_assignment/algorithm.cc | 215 ++++++- .../allocation_value.cc | 20 +- .../allocation_value.h | 7 + .../memory_space_assignment.proto | 36 ++ .../memory_space_assignment_test.cc | 527 +++++++++++++++++- .../memory_space_assignment_test_base.h | 6 + .../memory_space_assignment/options.cc | 2 + .../service/memory_space_assignment/options.h | 4 + .../prefetch_interval_picker.cc | 35 +- .../prefetch_interval_picker.h | 11 +- .../service/memory_space_assignment/utils.cc | 179 +++++- .../service/memory_space_assignment/utils.h | 44 +- 13 files changed, 1035 insertions(+), 56 deletions(-) diff --git a/third_party/xla/xla/service/memory_space_assignment/BUILD b/third_party/xla/xla/service/memory_space_assignment/BUILD index f0aa69e553c130..8595149933f3be 100644 --- a/third_party/xla/xla/service/memory_space_assignment/BUILD +++ b/third_party/xla/xla/service/memory_space_assignment/BUILD @@ -105,7 +105,6 @@ xla_test( ":utils", "//xla:comparison_util", "//xla:literal_util", - "//xla:shape_tree", "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", @@ -113,6 +112,7 @@ xla_test( "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:verified_hlo_module", + "//xla/hlo/transforms/simplifiers:instruction_hoister", "//xla/hlo/utils:hlo_live_range", "//xla/hlo/utils:hlo_matchers", "//xla/service:hlo_buffer", @@ -124,10 +124,8 @@ xla_test( "//xla/service/heap_simulator:allocation_block", "//xla/tests:test_utils", "//xla/tests:xla_internal_test_main", - "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:errors", "//xla/tsl/platform:logging", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -236,7 +234,6 @@ cc_library( "@highwayhash", "@highwayhash//:arch_specific", "@highwayhash//:hh_types", - "@tsl//tsl/platform:statusor", ], ) diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc index abf40b148bf0f9..f01572da6645f8 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc @@ -2210,9 +2210,13 @@ void MsaAlgorithm::CreateAllocationValuesForJointProcessedValues( defining_instruction->users().end(), may_be_replaced_by_slice_fn); if (!may_be_replaced_by_slice) { - VLOG(3) << "Skip " << interval.buffer->ToShortString() - << " because the buffer is larger than the heap size."; - continue; + if (!MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, + interval.buffer->defining_position(), interval.size)) { + VLOG(3) << "Skip " << interval.buffer->ToShortString() + << " because the buffer is larger than the heap size."; + continue; + } } } @@ -5412,6 +5416,11 @@ void MsaAlgorithm::MaybeSplitAllocationValues( bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation( AllocationValue& allocation_value) const { + if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { + return true; + } return allocation_value.value()->shape().has_layout() && allocation_value.value()->shape().layout().memory_space() == options_.alternate_memory_space; @@ -5420,7 +5429,10 @@ bool MsaAlgorithm::RequiresNoCopyAlternateMemAllocation( void MsaAlgorithm::AssignDefaultMemIfNotAllowedInAlternateMem( AllocationValue& allocation_value, int64_t definition_time) { if (!options_.is_position_allowed_in_alternate_mem_fn( - allocation_value.defining_position())) { + allocation_value.defining_position()) || + MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { std::optional existing_req = RequiredMemoryAssignmentAt(allocation_value.value(), definition_time); // If a value is pre-colored or already possesses an explicit alternate @@ -5697,6 +5709,15 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( << RequiresNoCopyAlternateMemAllocation(allocation_value); if (RequiresNoCopyAlternateMemAllocation(allocation_value) && allocation_value.size() > available_heap_size()) { + if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, + allocation_value.defining_position(), allocation_value.size())) { + return absl::ResourceExhaustedError(absl::StrCat( + "Cannot allocate pinned tensor in alternate memory: tensor size (", + allocation_value.size(), " bytes) exceeds available heap size (", + available_heap_size(), " bytes) for defining instruction '", + allocation_value.defining_instruction()->name(), "'")); + } VLOG(3) << "Skip " << allocation_value.value()->ToShortString() << " because the buffer is larger than the heap size."; continue; @@ -5830,6 +5851,79 @@ absl::StatusOr MsaAlgorithm::AllocateAllocationValues( request, result, options_.prefetch_interval_picker->retry_number()); } + + if (request.fail_on_unsatisfied_override || request.strict_timing || + (request.require_no_copy_alternate_mem_allocation && + MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, + allocation_value_to_update.defining_position(), + allocation_value_to_update.size()))) { + if (allocate_segment_result != AllocationResult::kSuccess) { + std::string reason = ResultToString(allocate_segment_result); + if (result_is(allocate_segment_result, + AllocationResult::kFailOutOfMemory)) { + reason = "Out of alternate memory / spatial placement conflict"; + } else if (result_is( + allocate_segment_result, + AllocationResult::kFailViolatesAsyncCopyResource)) { + reason = "Copy resource limit reached / bandwidth saturation"; + } else if (result_is(allocate_segment_result, + AllocationResult::kFailOutOfAsyncCopies)) { + reason = "Ran out of outstanding asynchronous copies"; + } else if (result_is(allocate_segment_result, + AllocationResult::kFailLiveRangeTooShort)) { + reason = "Live range too short / dependency ordering"; + } + return absl::FailedPreconditionError(absl::StrCat( + "MSA strict override failed for instruction '", + request.use->hlo_use.instruction->name(), "' operand ", + request.use->hlo_use.operand_number, ", shape ", + request.use->hlo_use.instruction + ->operand(request.use->hlo_use.operand_number) + ->shape() + .ToString(), + ", defining instruction '", + allocation_value_to_update.defining_instruction()->name(), + "': ", reason, " (requested schedule time: ", + request.preferred_prefetch_time.has_value() + ? absl::StrCat(*request.preferred_prefetch_time) + : "n/a", + ")")); + } + if (allocate_segment_result == AllocationResult::kSuccess && + request.strict_timing && + request.preferred_prefetch_time.has_value()) { + const Allocation* allocation = + allocation_sequence->empty() + ? nullptr + : allocation_sequence->back().get(); + int64_t actual_prefetch_time = 0; + if (allocation && allocation->is_copy_allocation()) { + actual_prefetch_time = + static_cast(allocation) + ->copy_start_schedule_after(); + } else if (allocation && allocation->is_sliced_copy_allocation()) { + actual_prefetch_time = + static_cast(allocation) + ->slice_details_sorted_by_start_time() + .front() + .copy_start_after_time; + } + if (actual_prefetch_time != *request.preferred_prefetch_time) { + return absl::FailedPreconditionError(absl::StrCat( + "MSA strict timing override unsatisfied for instruction '", + request.use->hlo_use.instruction->name(), "' operand ", + request.use->hlo_use.operand_number, ", shape ", + request.use->hlo_use.instruction + ->operand(request.use->hlo_use.operand_number) + ->shape() + .ToString(), + ": scheduled prefetch time (", actual_prefetch_time, + ") does not match requested prefetch time (", + *request.preferred_prefetch_time, ")")); + } + } + } if (allocate_segment_result == AllocationResult::kSuccess && NeedsMirroredAllocation(allocation_value_to_update, use, previous_use)) { @@ -6256,28 +6350,61 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( } } + bool strict_timing = false; + bool fail_on_unsatisfied_override = false; + bool is_prefetch_override = false; + bool require_end_colored_in_default_memory = false; + + if (MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + options_.msa_tensor_overrides, hlo_use, allocation_value.size()) || + MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { + allow_prefetch = false; + allow_no_copy_alternate_mem_allocation = false; + require_end_colored_in_default_memory = true; + AddRequiredAssignment( + allocation_value_to_update.value(), hlo_use.instruction, + MemorySpace::kDefault, use_time, + RequiredMemoryAssignment::Source::kUseNotAllowedInAlternateMemory); + } + if (MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + options_.msa_tensor_overrides, allocation_value.defining_position(), + allocation_value.size())) { + require_no_copy_alternate_mem_allocation = true; + allow_prefetch = false; + } + int64_t live_range_start_time = (earliest_prefetch_time.has_value() ? earliest_prefetch_time.value() : std::min(definition_time, use_time)); - auto overridden_preferred_prefetch_time = - MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( + auto prefetch_override_info = + MemorySpaceAssignmentUtils::GetPrefetchOverrideInfo( + options_.msa_tensor_overrides, options_.preferred_prefetch_overrides, allocation_value.size(), hlo_use, instruction_schedule, live_range_start_time, - latest_prefetch_time); - CHECK_OK(overridden_preferred_prefetch_time.status()); - if (overridden_preferred_prefetch_time.value().has_value()) { - VLOG(1) << "Overriding preferred prefetch for " - << hlo_use.instruction->name() << " operand number " - << hlo_use.operand_number << " operand index " - << hlo_use.operand_index.ToString() << " size " - << allocation_value.size() << " live range (" + latest_prefetch_time, allocation_value.defining_position()); + CHECK_OK(prefetch_override_info.status()); + if (prefetch_override_info.value().has_value()) { + VLOG(1) << "Overriding prefetch for " << hlo_use.instruction->name() + << " operand number " << hlo_use.operand_number + << " operand index " << hlo_use.operand_index.ToString() + << " size " << allocation_value.size() << " live range (" << live_range_start_time << ", " << latest_prefetch_time << ") from " << (preferred_prefetch_time.has_value() ? preferred_prefetch_time.value() : -1) - << " to " << overridden_preferred_prefetch_time.value().value(); - preferred_prefetch_time = overridden_preferred_prefetch_time.value(); + << " to " + << (prefetch_override_info->value().prefetch_time.has_value() + ? absl::StrCat( + *prefetch_override_info->value().prefetch_time) + : "none"); + preferred_prefetch_time = prefetch_override_info->value().prefetch_time; + strict_timing = prefetch_override_info->value().strict_timing; + fail_on_unsatisfied_override = + prefetch_override_info->value().fail_on_unsatisfied_override; + is_prefetch_override = true; } // Rarely, (e.g., when conditional true and false parameters are the @@ -6306,6 +6433,11 @@ AllocationRequest MsaAlgorithm::CreateAllocationRequest( request.required_copy_allocation_for = required_copy_allocation_for; request.required_copy_for_slice = required_copy_for_slice; request.allocation_value_to_update = &allocation_value_to_update; + request.strict_timing = strict_timing; + request.fail_on_unsatisfied_override = fail_on_unsatisfied_override; + request.is_prefetch_override = is_prefetch_override; + request.require_end_colored_in_default_memory = + require_end_colored_in_default_memory; } if (shape_override.has_value()) { @@ -8627,7 +8759,13 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { request.inclusive_start_time); std::optional required_memory_space_at_start; if (required_assignment_at_start.has_value()) { - required_memory_space_at_start = required_assignment_at_start->memory_space; + if (request.require_no_copy_alternate_mem_allocation && + required_assignment_at_start->memory_space == MemorySpace::kDefault) { + required_assignment_at_start = std::nullopt; + } else { + required_memory_space_at_start = + required_assignment_at_start->memory_space; + } } // Find required assignment both for the use and its aliases. If they are both // non-nullopt, then make sure they require the same assignment. @@ -8736,8 +8874,9 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { // First try keeping the allocation entirely in the alternate memory. if (!request.require_start_colored_in_default_memory && !request.require_end_colored_in_default_memory && - required_memory_space_at_start != MemorySpace::kDefault && - required_memory_space_at_end != MemorySpace::kDefault && + (request.require_no_copy_alternate_mem_allocation || + (required_memory_space_at_start != MemorySpace::kDefault && + required_memory_space_at_end != MemorySpace::kDefault)) && request.allow_no_copy_alternate_mem_allocation && !request.require_copy_allocation) { CheckAndUpdateForDualLiveAllocationValues(required_assignment_at_start, @@ -8752,7 +8891,11 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { } } - CHECK(!request.require_no_copy_alternate_mem_allocation); + if (request.require_no_copy_alternate_mem_allocation) { + return allocation_result != AllocationResult::kSuccess + ? allocation_result + : AllocationResult::kFailOutOfMemory; + } if (request.require_start_colored_in_alternate_memory) { // Since no-copy-allocation failed, continuous allocation is not possible in @@ -8916,6 +9059,9 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { << ") doesn't match the preferred prefetch time (" << *request.preferred_prefetch_time << "): " << request.use->hlo_use.ToString(); + if (request.strict_timing) { + return AllocationResult::kFailLiveRangeTooShort; + } } } return AllocationResult::kSuccess; @@ -8928,6 +9074,9 @@ AllocationResult MsaAlgorithm::AllocateSegment(AllocationRequest& request) { << ") which could not be satisfied: " << request.use->hlo_use.ToString(); } + if (request.fail_on_unsatisfied_override || request.strict_timing) { + return prefetch_result; + } result_mark(prefetch_result, allocation_result); } @@ -9177,9 +9326,11 @@ AllocationResult MsaAlgorithm::AllocateInAlternateMemoryNoCopy( bool can_eliminate_copy = false; if (request.allocation_value->allocation_sequence()->empty()) { // There hasn't been any allocations for this interval so far. We can - // eliminate copy if the value can be placed in the alternate memory. - can_eliminate_copy = options_.is_allowed_in_alternate_mem_fn( - *request.allocation_value->value()); + // eliminate copy if the value can be placed in the alternate memory or is + // pinned. + can_eliminate_copy = request.require_no_copy_alternate_mem_allocation || + options_.is_allowed_in_alternate_mem_fn( + *request.allocation_value->value()); } else { // If there has been a previous allocation, we can eliminate the copy if the // previous allocation was also in the alternate memory. @@ -10115,12 +10266,24 @@ AllocationResult MsaAlgorithm::InitializePrefetchIntervalPicker( std::optional preferred_prefetch_time = context.request->preferred_prefetch_time; if (preferred_prefetch_time) { - preferred_prefetch_time = - std::max(*preferred_prefetch_time, earliest_exclusive_prefetch_time); + if (context.request->strict_timing) { + if (*preferred_prefetch_time < earliest_exclusive_prefetch_time || + *preferred_prefetch_time >= context.prefetch_end_time) { + VLOG(3) << "Strict prefetch time " << *preferred_prefetch_time + << " is outside allowable prefetch range (" + << earliest_exclusive_prefetch_time << ", " + << context.prefetch_end_time << ")."; + return AllocationResult::kFailLiveRangeTooShort; + } + } else { + preferred_prefetch_time = + std::max(*preferred_prefetch_time, earliest_exclusive_prefetch_time); + } } options_.prefetch_interval_picker->Begin( context.request->use->hlo_use, earliest_exclusive_prefetch_time, - context.prefetch_end_time, preferred_prefetch_time); + context.prefetch_end_time, preferred_prefetch_time, + context.request->strict_timing); VLOG(3) << "Trying prefetch picker = " << options_.prefetch_interval_picker->ToDebugString(); diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc index bb372b08055c6e..d1df3be48dccde 100644 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_value.cc @@ -94,14 +94,18 @@ std::string AllocationRequest::ToString() const { no_copy_chunk_inclusive_start_time.has_value() ? absl::StrCat(*no_copy_chunk_inclusive_start_time) : "nullopt"), - absl::StrCat("require_start_colored_in_alternate_memmory: ", - require_start_colored_in_alternate_memory, - "; require_end_colored_in_alternate_memory: ", - require_end_colored_in_alternate_memory, - "; require_start_colored_in_default_memory: ", - require_start_colored_in_default_memory, - "; require_end_colored_in_default_memory: ", - require_end_colored_in_default_memory)}, + absl::StrCat( + "require_start_colored_in_alternate_memmory: ", + require_start_colored_in_alternate_memory, + "; require_end_colored_in_alternate_memory: ", + require_end_colored_in_alternate_memory, + "; require_start_colored_in_default_memory: ", + require_start_colored_in_default_memory, + "; require_end_colored_in_default_memory: ", + require_end_colored_in_default_memory, + "; strict_timing: ", strict_timing, + "; fail_on_unsatisfied_override: ", fail_on_unsatisfied_override, + "; is_prefetch_override: ", is_prefetch_override)}, "\n"); } diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h index ce623e08eb4c01..502335244f6b37 100644 --- a/third_party/xla/xla/service/memory_space_assignment/allocation_value.h +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_value.h @@ -291,6 +291,13 @@ struct AllocationRequest { // Indicates if the AllocationRequest end time (use time) has a default // memory color requirement. bool require_end_colored_in_default_memory = false; + // If true, forces the prefetch to start at the exact preferred_prefetch_time + // without searching outwards. + bool strict_timing = false; + // If true, any failure to satisfy this override results in a fatal error. + bool fail_on_unsatisfied_override = false; + // Indicates if this request has a prefetch override. + bool is_prefetch_override = false; std::string ToString() const; }; diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto index 0f1ed29b50357a..db0f1949eacfe9 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.proto @@ -113,7 +113,13 @@ message PreferredPrefetchOverrideOptions { // Preferred prefetch time is set to before the first instruction that // matches the filter. HloPositionMatcher before_instruction = 5; + // Preferred prefetch time set directly to a logical schedule time. + int64 logical_time = 6; } + // If true, the picker and allocator must ONLY attempt this exact time (no + // outward search). If it cannot be scheduled at this exact time, fail + // compilation. + optional bool strict_timing = 7; } // Filters operands in an HLO schedule and overrides preferred prefetch times @@ -123,12 +129,16 @@ message PreferredPrefetchOverride { optional HloOperandFilter hlo_operand_filter = 1; optional xla.memory_space_assignment.PreferredPrefetchOverrideOptions override_options = 2; + // If true, any failure to fulfill the action results in a fatal error. + optional bool fail_on_unsatisfied_override = 3; } // Encloses chained override configs. The first config has highest precedence // and so on. message PreferredPrefetchOverrides { repeated PreferredPrefetchOverride overrides = 1; + // Global flag to fail compilation if any override is unsatisfied. + optional bool fail_on_unsatisfied_override = 2; } // Specifies details on how to randomally select HloPositions for perturbation. @@ -213,6 +223,32 @@ message MsaSortOrderOverrides { repeated MsaSortOrderOverride overrides = 1; } +// Action to pin a buffer in alternate memory across its entire lifetime. +message PinInAlternateMemoryAction {} + +// Action to keep a buffer or operand use strictly in default memory (HBM). +message KeepInDefaultMemoryAction {} + +// Unified override rule for tensor and operand placement and timing. Overrides +// specified via MsaTensorOverride are always strictly enforced; compilation +// fails if an override cannot be fulfilled. +message MsaTensorOverride { + // Target filter for operand use or buffer position. + optional HloOperandFilter hlo_operand_filter = 1; + optional HloPositionMatcher hlo_position_matcher = 2; + + oneof action { + PreferredPrefetchOverrideOptions prefetch = 3; + PinInAlternateMemoryAction pin_in_alternate_memory = 4; + KeepInDefaultMemoryAction keep_in_default_memory = 5; + } +} + +// Encloses chained tensor override configs. +message MsaTensorOverrides { + repeated MsaTensorOverride overrides = 1; +} + // Expanded scoped alternate memory is a feature used at the end of MSA, in // in which we attempt to expand the size of allocated scoped alternate memory // buffers to the largest contiguous open space available. diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc index 5664b984154742..cdd421527b137a 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test.cc @@ -55,6 +55,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/hlo/transforms/simplifiers/instruction_hoister.h" #include "xla/hlo/utils/hlo_live_range.h" #include "xla/hlo/utils/hlo_matchers.h" #include "xla/layout.h" @@ -82,13 +83,10 @@ limitations under the License. #include "xla/service/memory_space_assignment/testing_utils.h" #include "xla/service/memory_space_assignment/utils.h" #include "xla/shape.h" -#include "xla/shape_tree.h" #include "xla/shape_util.h" #include "xla/tests/test_utils.h" -#include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/logging.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" #include "tsl/platform/protobuf.h" // IWYU pragma: keep @@ -2436,6 +2434,529 @@ TEST_F(MemorySpaceAssignmentTest, FilterUpdatePreferredPrefetchNoMatchTest) { EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); } +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchExactTimeSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { + after_instruction { instruction_name_regex: ".*negate.3.*" } + strict_timing: true + } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, + kDefaultMemorySpace, + op::Parameter(1)))); + // Ensure the CopyStart is scheduled after negate.3 (schedule index 6). + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[0], op::Parameter(0)); + EXPECT_THAT(sequence.instructions()[1], op::Parameter(1)); + EXPECT_THAT(sequence.instructions()[5], op::Negate()); + EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); + EXPECT_THAT(sequence.instructions()[10], op::CopyDone()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchLogicalTimeSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { logical_time: 4 strict_timing: true } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchUnsatisfiableFails) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate1, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { logical_time: 0 strict_timing: true } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + InstructionHoister instruction_hoister; + CHECK_OK(instruction_hoister.Run(module.get()).status()); + InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); + auto status_or = AssignMemorySpaceAndReturnStatus( + module.get(), std::move(options), /*buffer_interval_compare=*/{}, + &prefetch_interval_picker); + EXPECT_FALSE(status_or.ok()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPinInAlternateMemory) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate1, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*p1.*" } + pin_in_alternate_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + // p1 is pinned in alternate memory, so no copy allocation is created and + // add consumes p1 directly. + EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); + EXPECT_EQ(p1->shape().layout().memory_space(), kAlternateMemorySpace); +} + +TEST_F(MemorySpaceAssignmentTest, StrictKeepInDefaultMemory) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + // Strictly keep p1 in default memory, preventing prefetch. + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + keep_in_default_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + // Without override, p1 would be prefetched via AsyncCopy. + // With keep_in_default_memory, add consumes Parameter(1) directly. + EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); + EXPECT_EQ(p1->shape().layout().memory_space(), kDefaultMemorySpace); +} + +TEST_F(MemorySpaceAssignmentTest, StrictKeepInDefaultMemoryDefiningPosition) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*p1.*" } + keep_in_default_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + EXPECT_THAT(add, op::Add(op::Negate(), op::Parameter(1))); + EXPECT_EQ(p1->shape().layout().memory_space(), kDefaultMemorySpace); +} + +TEST_F(MemorySpaceAssignmentTest, + PositionMatcherDoesNotMatchConsumerInstruction) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + // Matching ".*add.*" with HloPositionMatcher should ONLY match buffers + // PRODUCED by add, NOT buffers CONSUMED by add (like p1). + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*add.*" } + keep_in_default_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + // Since p1's defining position is p1 (not add), it is NOT kept in default + // memory and gets prefetched into alternate memory space. + EXPECT_THAT(add, op::Add(op::Negate(), op::AsyncCopy(kAlternateMemorySpace, + kDefaultMemorySpace, + op::Parameter(1)))); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchBeforeInstructionSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { + before_instruction { instruction_name_regex: ".*negate.4.*" } + strict_timing: true + } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchExactTimeWithStrictBool) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { logical_time: 4 strict_timing: true } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpace(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[5], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPrefetchCostAnalysisExactTimeSuccess) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + HloInstruction* negate1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate0)); + HloInstruction* negate2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate1)); + HloInstruction* negate3 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate2)); + HloInstruction* negate4 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate3)); + HloInstruction* negate5 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate4)); + HloInstruction* negate6 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, negate5)); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate6, p1)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, p1, negate0, negate1, negate2, + negate3, negate4, negate5, negate6, add}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + + const std::string text_proto = R"pb( + overrides { + hlo_operand_filter { instruction_name_regex: ".*add.*" operand_number: 1 } + prefetch { + after_instruction { instruction_name_regex: ".*negate.3.*" } + strict_timing: true + } + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + AssignMemorySpaceUsingCostAnalysis(module.get(), std::move(options)); + + const HloInstructionSequence& sequence = + module->schedule().sequence(computation); + EXPECT_THAT(sequence.instructions()[6], op::CopyStart()); +} + +TEST_F(MemorySpaceAssignmentTest, StrictPinInAlternateMemoryFailsOnOom) { + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); // 24 bytes + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* negate0 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, p0)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + HloSchedule schedule(module.get()); + schedule.set_sequence(computation, {p0, negate0}); + CHECK_OK(module->set_schedule(schedule)); + + Options options = DefaultMemorySpaceOptions(); + options.max_size_in_bytes = 10; + const std::string text_proto = R"pb( + overrides { + hlo_position_matcher { instruction_name_regex: ".*p0.*" } + pin_in_alternate_memory {} + })pb"; + ASSERT_OK_AND_ASSIGN(options.msa_tensor_overrides, + ParseTextProto(text_proto)); + + InstructionHoister instruction_hoister; + CHECK_OK(instruction_hoister.Run(module.get()).status()); + InstructionCountPrefetchIntervalPicker prefetch_interval_picker(2, 10); + auto status_or = AssignMemorySpaceAndReturnStatus( + module.get(), std::move(options), /*buffer_interval_compare=*/{}, + &prefetch_interval_picker); + EXPECT_FALSE(status_or.ok()); +} + TEST_F(MemorySpaceAssignmentTest, EvictAndPrefetch) { std::unique_ptr module = CreateEvictAndPrefetchModule(); diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h index 8f869d9345e346..936607a2d2dc0b 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment_test_base.h @@ -304,6 +304,12 @@ class MemorySpaceAssignmentTestBase : public HloTestBase { if (options_override) { options = *std::move(options_override); } + for (const auto& override : options.msa_tensor_overrides.overrides()) { + if (override.has_pin_in_alternate_memory()) { + check_parameters_in_default_memory = false; + break; + } + } std::unique_ptr test_comparator; if (buffer_interval_compare.has_value()) { test_comparator = std::make_unique( diff --git a/third_party/xla/xla/service/memory_space_assignment/options.cc b/third_party/xla/xla/service/memory_space_assignment/options.cc index 8bbab64a6878ad..5bc554c0574cef 100644 --- a/third_party/xla/xla/service/memory_space_assignment/options.cc +++ b/third_party/xla/xla/service/memory_space_assignment/options.cc @@ -105,6 +105,8 @@ std::string Options::ToString() const { autotuning_config.has_value() ? "present" : "nullopt"), absl::StrCat("preferred_prefetch_overrides: \n", preferred_prefetch_overrides.DebugString()), + absl::StrCat("msa_tensor_overrides: \n", + msa_tensor_overrides.DebugString()), absl::StrCat("sliced_prefetch_options: \n", sliced_prefetch_options.DebugString()), absl::StrCat("memory_bound_loop_optimizer_options: \n", diff --git a/third_party/xla/xla/service/memory_space_assignment/options.h b/third_party/xla/xla/service/memory_space_assignment/options.h index b2635a1019bcc8..47fb95cc12ef17 100644 --- a/third_party/xla/xla/service/memory_space_assignment/options.h +++ b/third_party/xla/xla/service/memory_space_assignment/options.h @@ -407,6 +407,10 @@ struct Options { // filtered prefetches. PreferredPrefetchOverrides preferred_prefetch_overrides; + // Unified overrides for tensor placement, pinning, retention, and prefetch + // timing. + MsaTensorOverrides msa_tensor_overrides; + // Options for slicing prefetches into smaller asynchronously copied pieces. SlicedPrefetchOptions sliced_prefetch_options; diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc index d0594b1e9eace1..4a38f949752eed 100644 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc +++ b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.cc @@ -93,15 +93,22 @@ float InstructionCountPrefetchIntervalPicker::GetLogicalIntervalElapsed( void InstructionCountPrefetchIntervalPicker::Begin( const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) { + std::optional preferred_time, bool strict_timing) { end_time_ = end_time; + strict_timing_ = strict_timing; const Shape& shape = ShapeUtil::GetSubshape( use.instruction->operand(use.operand_number)->shape(), use.operand_index); if (preferred_time) { current_prefetch_time_ = *preferred_time; + if (strict_timing) { + strict_prefetch_time_ = *preferred_time; + } } else { current_prefetch_time_ = PreferredPrefetchStartTime(shape, start_time, end_time, end_time); + if (strict_timing) { + strict_prefetch_time_ = current_prefetch_time_; + } } } @@ -112,10 +119,16 @@ int64_t InstructionCountPrefetchIntervalPicker::Next() { } bool InstructionCountPrefetchIntervalPicker::Done() const { + if (strict_timing_) { + return current_prefetch_time_ > strict_prefetch_time_; + } return end_time_ - current_prefetch_time_ <= min_overlap_count_; } int64_t InstructionCountPrefetchIntervalPicker::latest_time() const { + if (strict_timing_) { + return strict_prefetch_time_; + } return end_time_ - min_overlap_count_ - 1; } @@ -356,7 +369,7 @@ int64_t CostAnalysisPrefetchIntervalPicker::EstimatedPrefetchEndTime( void CostAnalysisPrefetchIntervalPicker::Begin( const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) { + std::optional preferred_time, bool strict_timing) { const Shape& shape = ShapeUtil::GetSubshape( use.instruction->operand(use.operand_number)->shape(), use.operand_index); int64_t shape_size = size_override_ ? *size_override_ @@ -374,6 +387,17 @@ void CostAnalysisPrefetchIntervalPicker::Begin( end_logical_time_ = end_time; int end_nest_level = computation_nest_level_[end_logical_time_]; + if (strict_timing && preferred_time.has_value()) { + int64_t target_time = *preferred_time; + earliest_prefetch_time_ = target_time; + latest_prefetch_time_ = target_time; + increasing_prefetch_time_iterator_ = target_time; + decreasing_prefetch_time_iterator_ = target_time; + using_increasing_prefetch_time_iterator_ = true; + Next(); + return; + } + // Find the latest time we're allowed to start prefetching. float min_interval = min_overlap_to_async_copy_ratio_ * async_copy_elapsed_; latest_prefetch_time_ = @@ -527,8 +551,11 @@ std::string CostAnalysisPrefetchIntervalPicker::ToDebugString() const { int current_logical_prefetch_time = using_increasing_prefetch_time_iterator_ ? increasing_prefetch_time_iterator_ : decreasing_prefetch_time_iterator_; - float logical_interval_elapsed = GetLogicalIntervalElapsed( - current_logical_prefetch_time, end_logical_time_); + float logical_interval_elapsed = 0.0f; + if (current_logical_prefetch_time <= end_logical_time_) { + logical_interval_elapsed = GetLogicalIntervalElapsed( + current_logical_prefetch_time, end_logical_time_); + } return absl::StrCat( "Async copy elapsed (s) = ", async_copy_elapsed_, ", inst elapsed reduction (s) = ", inst_elapsed_reduction_, diff --git a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h index e12ab1a729da89..56eb842d1e8ebe 100644 --- a/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h +++ b/third_party/xla/xla/service/memory_space_assignment/prefetch_interval_picker.h @@ -83,7 +83,8 @@ class PrefetchIntervalPicker { // Begins the iterator for the first start time of the prefetch. virtual void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) = 0; + std::optional preferred_time, + bool strict_timing = false) = 0; // Advances the start time of the prefetch and returns that value. virtual int64_t Next() = 0; @@ -165,7 +166,8 @@ class InstructionCountPrefetchIntervalPicker : public PrefetchIntervalPicker { int64_t end_time) const override; void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) override; + std::optional preferred_time, + bool strict_timing = false) override; int64_t Next() override; bool Done() const override; @@ -181,6 +183,8 @@ class InstructionCountPrefetchIntervalPicker : public PrefetchIntervalPicker { int64_t max_overlap_count_; int64_t end_time_; int64_t current_prefetch_time_; + bool strict_timing_ = false; + int64_t strict_prefetch_time_ = -1; }; // Prefetch interval picker that uses cost analysis to overlap asynchronous @@ -231,7 +235,8 @@ class CostAnalysisPrefetchIntervalPicker : public PrefetchIntervalPicker { int64_t end_time) const override; void Begin(const HloUse& use, int64_t start_time, int64_t end_time, - std::optional preferred_time) override; + std::optional preferred_time, + bool strict_timing = false) override; int64_t Next() override; bool Done() const override; diff --git a/third_party/xla/xla/service/memory_space_assignment/utils.cc b/third_party/xla/xla/service/memory_space_assignment/utils.cc index f6b932a08e4b5c..363c3ac4f03364 100644 --- a/third_party/xla/xla/service/memory_space_assignment/utils.cc +++ b/third_party/xla/xla/service/memory_space_assignment/utils.cc @@ -44,7 +44,6 @@ limitations under the License. #include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" #include "xla/shape_util.h" #include "xla/util.h" -#include "tsl/platform/statusor.h" namespace xla { namespace memory_space_assignment { @@ -138,7 +137,7 @@ bool MemorySpaceAssignmentUtils::IsIntervalAllowedInAlternateMemory( }); } -bool MemorySpaceAssignmentUtils::DoesUseMatchFilter( +bool MemorySpaceAssignmentUtils::DoesUseMatchOperandFilter( const HloOperandFilter& filter, const HloUse& hlo_use, int64_t operand_size) { // The order of checks is such that the most expensive checks are done last. @@ -264,6 +263,49 @@ bool MemorySpaceAssignmentUtils::DoesPositionMatchFilter( DoesInstructionMatchRandomFilter(filter, *instruction); } +bool MemorySpaceAssignmentUtils::DoesPositionMatchPositionFilter( + const HloPositionMatcher& filter, const HloPosition& position, + int64_t size) { + if (filter.has_size_gte() && filter.size_gte() > size) { + return false; + } + if (filter.has_size_lte() && filter.size_lte() < size) { + return false; + } + if (filter.has_tuple_index() && + position.index != ShapeIndex(filter.tuple_index().index().begin(), + filter.tuple_index().index().end())) { + return false; + } + return DoesInstructionMatchFilter(filter, *position.instruction) && + DoesInstructionMatchRandomFilter(filter, *position.instruction); +} + +bool MemorySpaceAssignmentUtils::DoesPositionMatchOperandFilter( + const HloOperandFilter& filter, const HloPosition& position, int64_t size) { + if (filter.has_size_gte() && filter.size_gte() > size) { + return false; + } + if (filter.has_size_lte() && filter.size_lte() < size) { + return false; + } + if (filter.has_tuple_index() && + position.index != ShapeIndex(filter.tuple_index().index().begin(), + filter.tuple_index().index().end())) { + return false; + } + if (filter.has_instruction_name_regex() && + !RE2::FullMatch(position.instruction->name(), + filter.instruction_name_regex())) { + return false; + } + if (filter.has_instruction_regex() && + !RE2::FullMatch(position.instruction->ToString(), + filter.instruction_regex())) { + return false; + } + return true; +} bool MemorySpaceAssignmentUtils::DoesInstructionMatchFilter( const HloPositionMatcher& filter, const HloInstruction& instruction) { if (filter.has_instruction_name_regex() && @@ -284,8 +326,8 @@ bool MemorySpaceAssignmentUtils::DoesBufferIntervalMatchHloUseFilter( return true; } for (const HloUse& use : buffer_interval.buffer->GetUses()) { - if (DoesUseMatchFilter(filter.hlo_use_filter(), use, - buffer_interval.size)) { + if (DoesUseMatchOperandFilter(filter.hlo_use_filter(), use, + buffer_interval.size)) { return true; } } @@ -358,6 +400,9 @@ MemorySpaceAssignmentUtils::GetPrefetchTime( case PreferredPrefetchOverrideOptions::kBeforeInstruction: return GetPrefetchTimeBeforeInstruction( override_options.before_instruction(), instruction_schedule); + case PreferredPrefetchOverrideOptions::kLogicalTime: + return static_cast>( + override_options.logical_time()); case PreferredPrefetchOverrideOptions::OPTIONS_NOT_SET: break; } @@ -372,7 +417,7 @@ MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( instruction_schedule, int64_t earliest_prefetch_time, int64_t latest_prefetch_time) { for (const auto& override : preferred_prefetch_overrides.overrides()) { - if (!MemorySpaceAssignmentUtils::DoesUseMatchFilter( + if (!MemorySpaceAssignmentUtils::DoesUseMatchOperandFilter( override.hlo_operand_filter(), hlo_use, operand_size)) { continue; } @@ -394,6 +439,130 @@ MemorySpaceAssignmentUtils::GetOverriddenPreferredPrefetchTime( return static_cast>>(std::nullopt); } +absl::StatusOr> +MemorySpaceAssignmentUtils::GetPrefetchOverrideInfo( + const MsaTensorOverrides& msa_tensor_overrides, + const PreferredPrefetchOverrides& preferred_prefetch_overrides, + int64_t operand_size, const HloUse& hlo_use, + const absl::flat_hash_map& + instruction_schedule, + int64_t earliest_prefetch_time, int64_t latest_prefetch_time, + const std::optional& position) { + // First check msa_tensor_overrides. + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_prefetch()) { + continue; + } + bool matches = false; + if (override.has_hlo_operand_filter() && + DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, + operand_size)) { + matches = true; + } else if (override.has_hlo_position_matcher() && position.has_value() && + DoesPositionMatchPositionFilter(override.hlo_position_matcher(), + *position, operand_size)) { + matches = true; + } + if (!matches) { + continue; + } + const auto& prefetch_options = override.prefetch(); + ABSL_ASSIGN_OR_RETURN( + auto prefetch_time, + GetPrefetchTime(prefetch_options, earliest_prefetch_time, + latest_prefetch_time, instruction_schedule)); + bool strict_timing = prefetch_options.has_strict_timing() && + prefetch_options.strict_timing(); + PrefetchOverrideInfo info; + info.prefetch_time = prefetch_time; + info.strict_timing = strict_timing; + info.fail_on_unsatisfied_override = true; + return info; + } + + // Next check preferred_prefetch_overrides. + for (const auto& override : preferred_prefetch_overrides.overrides()) { + if (!DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, + operand_size)) { + continue; + } + ABSL_ASSIGN_OR_RETURN( + auto prefetch_time, + GetPrefetchTime(override.override_options(), earliest_prefetch_time, + latest_prefetch_time, instruction_schedule)); + bool strict_timing = override.override_options().has_strict_timing() && + override.override_options().strict_timing(); + bool fail_on_unsatisfied = + override.fail_on_unsatisfied_override() || + preferred_prefetch_overrides.fail_on_unsatisfied_override() || + strict_timing; + PrefetchOverrideInfo info; + info.prefetch_time = prefetch_time; + info.strict_timing = strict_timing; + info.fail_on_unsatisfied_override = fail_on_unsatisfied; + return info; + } + return std::nullopt; +} + +bool MemorySpaceAssignmentUtils::ShouldPinInAlternateMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloPosition& position, + int64_t size) { + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_pin_in_alternate_memory()) { + continue; + } + if (override.has_hlo_position_matcher() && + DoesPositionMatchPositionFilter(override.hlo_position_matcher(), + position, size)) { + return true; + } + if (override.has_hlo_operand_filter() && + DoesPositionMatchOperandFilter(override.hlo_operand_filter(), position, + size)) { + return true; + } + } + return false; +} + +bool MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloPosition& position, + int64_t size) { + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_keep_in_default_memory()) { + continue; + } + if (override.has_hlo_position_matcher() && + DoesPositionMatchPositionFilter(override.hlo_position_matcher(), + position, size)) { + return true; + } + if (override.has_hlo_operand_filter() && + DoesPositionMatchOperandFilter(override.hlo_operand_filter(), position, + size)) { + return true; + } + } + return false; +} + +bool MemorySpaceAssignmentUtils::ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloUse& hlo_use, + int64_t operand_size) { + for (const auto& override : msa_tensor_overrides.overrides()) { + if (!override.has_keep_in_default_memory()) { + continue; + } + if (override.has_hlo_operand_filter() && + DoesUseMatchOperandFilter(override.hlo_operand_filter(), hlo_use, + operand_size)) { + return true; + } + } + return false; +} + bool MemorySpaceAssignmentUtils::DoesCrossProgramPrefetchBufferMatchAnyFilter( const MsaSortOrderOverrides& sort_order_overrides, const MsaBufferInterval& buffer_interval) { diff --git a/third_party/xla/xla/service/memory_space_assignment/utils.h b/third_party/xla/xla/service/memory_space_assignment/utils.h index a7bd6f4688dbb3..b8a897252f3e35 100644 --- a/third_party/xla/xla/service/memory_space_assignment/utils.h +++ b/third_party/xla/xla/service/memory_space_assignment/utils.h @@ -18,7 +18,6 @@ limitations under the License. #include #include -#include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -36,6 +35,13 @@ namespace memory_space_assignment { using MsaBufferInterval = GlobalDecreasingSizeBestFitHeap::BufferInterval; +// Info resolved from prefetch overrides. +struct PrefetchOverrideInfo { + std::optional prefetch_time; + bool strict_timing = false; + bool fail_on_unsatisfied_override = false; +}; + // Encapsulates common utility methods for memory space assignment. class MemorySpaceAssignmentUtils { public: @@ -55,8 +61,9 @@ class MemorySpaceAssignmentUtils { const HloInstruction& instruction, const absl::flat_hash_set& execution_threads); - static bool DoesUseMatchFilter(const HloOperandFilter& filter, - const HloUse& hlo_use, int64_t operand_size); + static bool DoesUseMatchOperandFilter(const HloOperandFilter& filter, + const HloUse& hlo_use, + int64_t operand_size); static bool DoesInstructionMatchFilter(const HloPositionMatcher& filter, const HloInstruction& instruction); @@ -67,6 +74,14 @@ class MemorySpaceAssignmentUtils { static bool DoesPositionMatchFilter(const HloPositionMatcher& filter, const MsaBufferInterval& buffer_interval); + static bool DoesPositionMatchPositionFilter(const HloPositionMatcher& filter, + const HloPosition& position, + int64_t size); + + static bool DoesPositionMatchOperandFilter(const HloOperandFilter& filter, + const HloPosition& position, + int64_t size); + static absl::StatusOr GetScheduleTimeFromInstructionMatcher( const HloPositionMatcher& position_matcher, @@ -104,6 +119,29 @@ class MemorySpaceAssignmentUtils { instruction_schedule, int64_t earliest_prefetch_time, int64_t latest_prefetch_time); + static absl::StatusOr> + GetPrefetchOverrideInfo( + const MsaTensorOverrides& msa_tensor_overrides, + const PreferredPrefetchOverrides& preferred_prefetch_overrides, + int64_t operand_size, const HloUse& hlo_use, + const absl::flat_hash_map& + instruction_schedule, + int64_t earliest_prefetch_time, int64_t latest_prefetch_time, + const std::optional& position = std::nullopt); + + static bool ShouldPinInAlternateMemory( + const MsaTensorOverrides& msa_tensor_overrides, + const HloPosition& position, int64_t size); + + static bool ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, + const HloPosition& position, int64_t size); + + static bool ShouldKeepInDefaultMemory( + const MsaTensorOverrides& msa_tensor_overrides, const HloUse& hlo_use, + int64_t operand_size); + static bool DoesCrossProgramPrefetchBufferMatchAnyFilter( const MsaSortOrderOverrides& sort_order_overrides, const MsaBufferInterval& buffer_interval); From 85d824abd4e6db929e074e2aab91441029b17aae Mon Sep 17 00:00:00 2001 From: Mikhail Goncharov Date: Fri, 28 Aug 2026 04:30:32 -0700 Subject: [PATCH 22/23] [XLA:GPU] increase timeout for triton_fusion_numerics_verifier_test PiperOrigin-RevId: 972505706 --- third_party/xla/xla/backends/gpu/transforms/BUILD | 1 - 1 file changed, 1 deletion(-) diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 5d521b072e4bc5..3d26b3168889e9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -3832,7 +3832,6 @@ cc_library( xla_test( name = "triton_fusion_numerics_verifier_test", - timeout = "short", srcs = ["triton_fusion_numerics_verifier_test.cc"], backends = [ "a100", From 97878acda5ab57b0f84ed65c832fb5be3fe3d692 Mon Sep 17 00:00:00 2001 From: Penporn Koanantakool Date: Fri, 28 Aug 2026 04:31:09 -0700 Subject: [PATCH 23/23] [XLA:CPU:Autotuner] Rename `LlvmKernelAutotuner` to `CpuAutotuner` Use a more generic name since we will add more non-LLVM backends. PiperOrigin-RevId: 972505937 --- .../xla/xla/backends/cpu/autotuner/BUILD | 12 ++++++------ ...vm_kernel_autotuner.cc => cpu_autotuner.cc} | 9 +++++---- ...llvm_kernel_autotuner.h => cpu_autotuner.h} | 18 ++++++++---------- ...autotuner_test.cc => cpu_autotuner_test.cc} | 8 ++++---- 4 files changed, 23 insertions(+), 24 deletions(-) rename third_party/xla/xla/backends/cpu/autotuner/{llvm_kernel_autotuner.cc => cpu_autotuner.cc} (92%) rename third_party/xla/xla/backends/cpu/autotuner/{llvm_kernel_autotuner.h => cpu_autotuner.h} (68%) rename third_party/xla/xla/backends/cpu/autotuner/{llvm_kernel_autotuner_test.cc => cpu_autotuner_test.cc} (86%) diff --git a/third_party/xla/xla/backends/cpu/autotuner/BUILD b/third_party/xla/xla/backends/cpu/autotuner/BUILD index f37dba0ec46e67..460ef37e8b3379 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/BUILD +++ b/third_party/xla/xla/backends/cpu/autotuner/BUILD @@ -142,9 +142,9 @@ xla_cc_test( ) cc_library( - name = "llvm_kernel_autotuner", - srcs = ["llvm_kernel_autotuner.cc"], - hdrs = ["llvm_kernel_autotuner.h"], + name = "cpu_autotuner", + srcs = ["cpu_autotuner.cc"], + hdrs = ["cpu_autotuner.h"], deps = [ ":cpu_codegen_backend", ":cpu_profiler", @@ -171,10 +171,10 @@ cc_library( ) xla_cc_test( - name = "llvm_kernel_autotuner_test", - srcs = ["llvm_kernel_autotuner_test.cc"], + name = "cpu_autotuner_test", + srcs = ["cpu_autotuner_test.cc"], deps = [ - ":llvm_kernel_autotuner", + ":cpu_autotuner", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/service/cpu:cpu_compiler", diff --git a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.cc b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.cc similarity index 92% rename from third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.cc rename to third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.cc index 4408a49f4e3ccc..ee7792d167deda 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.cc +++ b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/backends/cpu/autotuner/llvm_kernel_autotuner.h" +#include "xla/backends/cpu/autotuner/cpu_autotuner.h" #include #include @@ -42,15 +42,16 @@ limitations under the License. namespace xla::cpu { -absl::StatusOr LlvmKernelAutotuner::RunImpl( +absl::StatusOr CpuAutotuner::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { ABSL_ASSIGN_OR_RETURN(auto compiler, CpuCodegenBackend::CreateBackendCompiler()); - ABSL_ASSIGN_OR_RETURN(auto backend, LlvmKernelBackend::Create(compiler.get())); + ABSL_ASSIGN_OR_RETURN(auto llvm_kernel_backend, + LlvmKernelBackend::Create(compiler.get())); std::unique_ptr profiler = CpuProfiler::Create(ProfileOptions()); std::vector> codegen_backends; - codegen_backends.push_back(std::move(backend)); + codegen_backends.push_back(std::move(llvm_kernel_backend)); ABSL_ASSIGN_OR_RETURN(auto orchestrator, CodegenOrchestrator::Create(std::move(codegen_backends), diff --git a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.h b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.h similarity index 68% rename from third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.h rename to third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.h index fe26d971346742..c76882157e527c 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner.h +++ b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef XLA_BACKENDS_CPU_AUTOTUNER_LLVM_KERNEL_AUTOTUNER_H_ -#define XLA_BACKENDS_CPU_AUTOTUNER_LLVM_KERNEL_AUTOTUNER_H_ +#ifndef XLA_BACKENDS_CPU_AUTOTUNER_CPU_AUTOTUNER_H_ +#define XLA_BACKENDS_CPU_AUTOTUNER_CPU_AUTOTUNER_H_ #include "absl/container/flat_hash_set.h" #include "absl/log/check.h" @@ -29,16 +29,14 @@ limitations under the License. namespace xla::cpu { -inline constexpr absl::string_view kLlvmKernelAutotunerName = - "llvm_kernel_autotuner"; +inline constexpr absl::string_view kCpuAutotunerName = "cpu_autotuner"; -// Llvm kernel autotuning pass. It tries to autotune the llvm kernel compilation -// provided by the LlvmKernelBackend. -class LlvmKernelAutotuner : public HloModulePass { +// CPU autotuning pass. +class CpuAutotuner : public HloModulePass { public: - LlvmKernelAutotuner() = default; + CpuAutotuner() = default; - absl::string_view name() const override { return kLlvmKernelAutotunerName; } + absl::string_view name() const override { return kCpuAutotunerName; } protected: absl::StatusOr RunImpl( @@ -48,4 +46,4 @@ class LlvmKernelAutotuner : public HloModulePass { } // namespace xla::cpu -#endif // XLA_BACKENDS_CPU_AUTOTUNER_LLVM_KERNEL_AUTOTUNER_H_ +#endif // XLA_BACKENDS_CPU_AUTOTUNER_CPU_AUTOTUNER_H_ diff --git a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner_test.cc b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner_test.cc similarity index 86% rename from third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner_test.cc rename to third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner_test.cc index f0b0bee12b7ffe..450bdc89fea5c6 100644 --- a/third_party/xla/xla/backends/cpu/autotuner/llvm_kernel_autotuner_test.cc +++ b/third_party/xla/xla/backends/cpu/autotuner/cpu_autotuner_test.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "xla/backends/cpu/autotuner/llvm_kernel_autotuner.h" +#include "xla/backends/cpu/autotuner/cpu_autotuner.h" #include @@ -36,10 +36,10 @@ constexpr absl::string_view kLlvmKernelConcatenateHlo = R"( } )"; -class LlvmKernelAutotunerTest : public HloHardwareIndependentTestBase {}; +class CpuAutotunerTest : public HloHardwareIndependentTestBase {}; -TEST_F(LlvmKernelAutotunerTest, GetBestConfig) { - LlvmKernelAutotuner autotuner; +TEST_F(CpuAutotunerTest, GetBestConfig) { + CpuAutotuner autotuner; TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr module, ParseAndReturnVerifiedModule(kLlvmKernelConcatenateHlo));