[XLA:GPU] Rely on LLVM parser rather than objcopy to load fatbin in t…

…ests

To avoid relying on `objcopy` from toolchains

PiperOrigin-RevId: 712832511

xla/stream_executor/gpu/BUILD

@@ -552,33 +552,28 @@ gpu_kernel_library(
     ]),
 )
 
+# Extract the .so file from the gpu_test_kernels library.
+# TODO: make gpu_test_kernels a direct dependency of gpu_test_kernels_fatbin.
 genrule(
-    name = "gpu_test_kernels_fatbin_extractor",
+    name = "gpu_test_kernels_object_extractor",
     testonly = True,
     srcs = [":gpu_test_kernels"],
-    outs = ["gpu_test_kernels.fatbin"],
+    outs = ["gpu_test_kernels.so"],
     cmd = """
-      STATIC_LIBRARY=""
+      SHARED_OBJECT=""
       for src in $(SRCS); do
-        if [[ $$src == *.a ]]; then
-          STATIC_LIBRARY=$$src
-          break
-        fi
-      done
-
-      if [[ -z $$STATIC_LIBRARY ]]; then
-        echo "No static library found in $(SRCS)" >&2
-        exit 1
+      if [[ $$src == *.so ]]; then
+      SHARED_OBJECT=$$src
+      cp $$src $@ # Copy the .so file to the output
+      break
       fi
+      done
 
-      $(OBJCOPY) "--dump-section=.nv_fatbin=$@" "$$STATIC_LIBRARY" || true
-      
-      if [ ! -f "$@" ]; then
-        # binutils' objcopy doesn't return a non-zero exit code if the
-        # section was not found, so we need to check for the file's existence instead.
-        $(OBJCOPY) "--dump-section=.hip_fatbin=$@" "$$STATIC_LIBRARY"
+      if [[ -z $$SHARED_OBJECT ]]; then
+      echo "No .so file found in $(SRCS)" >&2
+      exit 1
       fi
-    """,
+      """,
     tags = ["gpu"],
     toolchains = ["@bazel_tools//tools/cpp:current_cc_toolchain"],
 )
@@ -588,17 +583,35 @@ cc_library(
     testonly = True,
     srcs = ["gpu_test_kernels_fatbin.cc"],
     hdrs = ["gpu_test_kernels_fatbin.h"],
-    data = [":gpu_test_kernels_fatbin_extractor"],
+    data = [
+        ":gpu_test_kernels_object_extractor",
+    ],
     tags = ["gpu"],
     deps = [
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Object",
+        "@llvm-project//llvm:Support",
         "@tsl//tsl/platform:env",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:path",
         "@tsl//tsl/platform:test",
     ],
 )
 
+xla_test(
+    name = "gpu_test_kernels_fatbin_test",
+    srcs = ["gpu_test_kernels_fatbin_test.cc"],
+    backends = ["gpu"],
+    deps = [
+        ":gpu_test_kernels_fatbin",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:test",
+        "@tsl//tsl/platform:test_main",
+    ],
+)
+
 xla_test(
     name = "gpu_kernel_test",
     srcs = ["gpu_kernel_test.cc"],

xla/stream_executor/gpu/gpu_test_kernels_fatbin.cc

@@ -19,7 +19,15 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/match.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/Object/ObjectFile.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/path.h"
@@ -31,9 +39,38 @@ absl::StatusOr<std::vector<uint8_t>> GetGpuTestKernelsFatbin() {
   tsl::Env* env = tsl::Env::Default();
   std::string file_path =
       tsl::io::JoinPath(tsl::testing::XlaSrcRoot(), "stream_executor", "gpu",
-                        "gpu_test_kernels.fatbin");
+                        "gpu_test_kernels.so");
+
   std::string file_contents;
   TF_RETURN_IF_ERROR(tsl::ReadFileToString(env, file_path, &file_contents));
-  return std::vector<uint8_t>(file_contents.begin(), file_contents.end());
+
+  const auto buffer = llvm::MemoryBuffer::getMemBuffer(
+      llvm::StringRef(file_contents),
+      /*BufferName=*/"", /*RequiresNullTerminator=*/false);
+  auto object_file =
+      llvm::object::ObjectFile::createObjectFile(buffer->getMemBufferRef());
+
+  if (!object_file) {
+    return absl::InternalError(llvm::toString(object_file.takeError()));
+  }
+
+  const auto executable_elf_object_file =
+      llvm::dyn_cast<llvm::object::ELF64LEObjectFile>(object_file.get().get());
+
+  if (!executable_elf_object_file) {
+    return absl::InternalError(
+        "Generated executable binary is not a 64bit ELF file.");
+  }
+
+  for (const auto& section : executable_elf_object_file->sections()) {
+    if (absl::StartsWith(section.getName().get().str(), ".nv_fatbin") ||
+        absl::StartsWith(section.getName().get().str(), ".hip_fatbin")) {
+      const std::string fatbin_contents = section.getContents().get().str();
+      return std::vector<uint8_t>(fatbin_contents.begin(),
+                                  fatbin_contents.end());
+    }
+  }
+
+  return absl::InternalError("Fatbin section not found in generated ELF file.");
 }
 }  // namespace stream_executor::gpu

xla/stream_executor/gpu/gpu_test_kernels_fatbin_test.cc

@@ -0,0 +1,35 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/stream_executor/gpu/gpu_test_kernels_fatbin.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/test.h"
+
+namespace stream_executor::gpu {
+namespace {
+
+TEST(GpuTestKernelsFatbinTest, GetGpuTestKernelsFatbin) {
+  std::vector<uint8_t> fatbin;
+
+  TF_ASSERT_OK_AND_ASSIGN(fatbin, GetGpuTestKernelsFatbin());
+  EXPECT_FALSE(fatbin.empty());
+}
+
+}  // namespace
+}  // namespace stream_executor::gpu