[XLA:GPU] update nvjitlink and compilation_provider tests to support cuda 12.8

Update test expectations due to cuda 12.8 changes nvJitLinkCreate behavior to fail when an invalid sm architecture is provided

Update nvJitLinkDestroy usage as asan detects a leakage if it is not run after nvJitLinkCreate failure.

PiperOrigin-RevId: 743222408

Author: IllogicalMoose

third_party/xla/xla/stream_executor/cuda/BUILD

@@ -930,6 +930,7 @@ xla_cc_test(
         ":nvjitlink_support",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/gpu:gpu_asm_opts",
+        "//xla/tsl/platform:status_matchers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",

third_party/xla/xla/stream_executor/cuda/nvjitlink_impl.cc

@@ -163,13 +163,25 @@ absl::StatusOr<std::vector<uint8_t>> CompileAndLinkUsingLibNvJitLink(
   absl::c_transform(cli_args, std::back_inserter(cli_args_ptrs),
                     [](const std::string& s) { return s.c_str(); });
 
-  nvJitLinkHandle link_handle{};
-  RETURN_IF_NVJITLINK_ERROR(nvJitLinkCreate(&link_handle, cli_args_ptrs.size(),
-                                            cli_args_ptrs.data()));
+  nvJitLinkHandle link_handle = nullptr;
+  nvJitLinkResult create_result =
+      nvJitLinkCreate(&link_handle, cli_args_ptrs.size(), cli_args_ptrs.data());
+
   absl::Cleanup link_handle_cleaner = [&link_handle] {
-    CHECK_EQ(nvJitLinkDestroy(&link_handle), NVJITLINK_SUCCESS);
+    if (link_handle != nullptr) {
+      CHECK_EQ(nvJitLinkDestroy(&link_handle), NVJITLINK_SUCCESS);
+    }
   };
 
+  if (create_result != NVJITLINK_SUCCESS) {
+    TF_ASSIGN_OR_RETURN(std::string error_log,
+                        nvJitLinkGetErrorLog(link_handle));
+
+    VLOG(3) << "libnvjitlink error log output: " << error_log;
+
+    return ToStatus(create_result, error_log);
+  }
+
   for (auto& image : inputs) {
     nvJitLinkInputType input_type = image.type == NvJitLinkInput::Type::kPtx
                                         ? NVJITLINK_INPUT_PTX

third_party/xla/xla/stream_executor/cuda/nvjitlink_test.cc

@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/nvjitlink_support.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/gpu_asm_opts.h"
+#include "xla/tsl/platform/status_matchers.h"
 #include "tsl/platform/status_matchers.h"
 #include "tsl/platform/test.h"
 
@@ -165,7 +166,8 @@ TEST_F(NvJitLinkTest, IdentifiesUnsupportedArchitecture) {
   EXPECT_THAT(
       CompileAndLinkHelper(stream_executor::CudaComputeCapability{100, 0},
                            {kStandalonePtx}),
-      tsl::testing::StatusIs(absl::StatusCode::kUnimplemented));
+      tsl::testing::StatusIs(testing::AnyOf(absl::StatusCode::kUnknown,
+                                            absl::StatusCode::kUnimplemented)));
 }
 
 TEST_F(NvJitLinkTest, LinkingTwoCompilationUnitsSucceeds) {