From ccdd3c14d124e17e976ad1550fd64e9204bfe634 Mon Sep 17 00:00:00 2001 From: Henning Becker Date: Tue, 17 Sep 2024 03:06:05 -0700 Subject: [PATCH] Fix ptx_compilation_test failure on H100 Unfortunately the test is still not as robust as I would like it to be. Some slightly differently generated PTX from Triton leads to some of the comparisons fails. In particular when comparing PTX compiled in one-go with PTX first compiled to a relocatable object and then linked into a binary. The solution for now is to not compare relocatable PTX compilation against non-relocatable PTX compilation. I'm also disabling autotuning as a precaution - even though it was not the cause of this issue. PiperOrigin-RevId: 675490844 --- xla/service/gpu/ptx_compilation_test.cc | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/xla/service/gpu/ptx_compilation_test.cc b/xla/service/gpu/ptx_compilation_test.cc index 1d194c06b2988..967c0e494d67c 100644 --- a/xla/service/gpu/ptx_compilation_test.cc +++ b/xla/service/gpu/ptx_compilation_test.cc @@ -217,6 +217,12 @@ class NVPTXCompilationTests debug_options->set_xla_llvm_force_inline_before_split(false); } + DebugOptions GetDebugOptionsForTest() override { + auto debug_options = HloTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_autotune_level(0); + return debug_options; + } + void SetUp() override { HloTestBase::SetUp(); std::string_view name = std::get<0>(GetParam()); @@ -283,10 +289,17 @@ TEST_P(NVPTXCompilationTests, CompareBinaryOutput) { absl::StatusOr> executable = compile(compilation_method, linking_method); - constexpr PtxLinkingMethod kReferenceLinkingMethod = - PtxLinkingMethod::kNvJitLink; + // Binaries produced in a separate linking step differ from binaries produced + // with combined compilation/linking. Therefore we only enable linking in the + // reference build when the build under test also uses a separate linking + // step. + const PtxLinkingMethod reference_linking_method = + (linking_method == PtxLinkingMethod::kNone) + ? PtxLinkingMethod::kNone + : PtxLinkingMethod::kNvJitLink; + absl::StatusOr> reference = - compile(PtxCompilationMethod::kPtxas, kReferenceLinkingMethod); + compile(PtxCompilationMethod::kPtxas, reference_linking_method); EXPECT_THAT(executable, tsl::testing::IsOkAndHolds(::testing::NotNull())); EXPECT_THAT(reference, tsl::testing::IsOkAndHolds(::testing::NotNull())); @@ -349,7 +362,7 @@ TEST_P(NVPTXCompilationTests, CompareBinaryOutput) { TF_ASSERT_OK_AND_ASSIGN(auto reference_text_sections, get_text_sections(reference_binary)); - if (linking_method == kReferenceLinkingMethod) { + if (linking_method == reference_linking_method) { EXPECT_THAT(executable_text_sections, ::testing::Eq(reference_text_sections)); return;