From ccdd3c14d124e17e976ad1550fd64e9204bfe634 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
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<std::unique_ptr<Executable>> 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<std::unique_ptr<Executable>> 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;