Skip to content

Commit

Permalink
Fix ptx_compilation_test failure on H100
Browse files Browse the repository at this point in the history
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
  • Loading branch information
beckerhe authored and Google-ML-Automation committed Sep 17, 2024
1 parent e51e376 commit ccdd3c1
Showing 1 changed file with 17 additions and 4 deletions.
21 changes: 17 additions & 4 deletions xla/service/gpu/ptx_compilation_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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()));
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit ccdd3c1

Please sign in to comment.