PR #21822: [XLA:GPU] Add support for SM100a architecture (Blackwell)

Imported from GitHub PR #21822

Created `ShouldUsePtxExtension` helper for the extension suffix (this will also be used for sm120, etc).

CUDA 12.8 was recently released, which supports PTX 8.7, but that is not supported by the integrated LLVM (support added in llvm/llvm-project#124155), so leaving the association with PTX 8.6 - this doesn't raise warnings during compilation.

Copybara import of the project:

--
267cf74 by Sergey Kozub <skozub@nvidia.com>:

Add support for SM100a architecture (Blackwell)

Merging this change closes #21822

COPYBARA_INTEGRATE_REVIEW=#21822 from openxla:devel/sm100a 267cf74
PiperOrigin-RevId: 720806648

Author: sergey-kozub

xla/service/gpu/llvm_gpu_backend/BUILD

@@ -99,6 +99,7 @@ cc_library(
         "//xla/service/llvm_ir:llvm_command_line_options",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:semantic_version",
+        "//xla/stream_executor/cuda:ptx_compiler_helpers",
         "//xla/stream_executor/cuda:subprocess_compilation",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/status",

xla/service/gpu/llvm_gpu_backend/nvptx_backend.cc

@@ -60,6 +60,7 @@ limitations under the License.
 #include "xla/service/gpu/llvm_gpu_backend/nvptx_libdevice_path.h"
 #include "xla/service/gpu/metrics.h"
 #include "xla/service/llvm_ir/llvm_command_line_options.h"
+#include "xla/stream_executor/cuda/ptx_compiler_helpers.h"
 #include "xla/stream_executor/cuda/subprocess_compilation.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/semantic_version.h"
@@ -237,8 +238,8 @@ std::string GetSmName(se::CudaComputeCapability compute_capability) {
   int sm_version = 30;
   // If the current compute capability isn't known, fallback to the
   // most recent version before it.
-  int supported_versions[] = {90, 89, 87, 86, 80, 75, 72, 70, 62,
-                              61, 60, 53, 52, 50, 37, 35, 32, 30};
+  int supported_versions[] = {100, 90, 89, 87, 86, 80, 75, 72, 70, 62,
+                              61,  60, 53, 52, 50, 37, 35, 32, 30};
   for (int v : supported_versions) {
     if (v <= compute_capability_version) {
       sm_version = v;
@@ -260,8 +261,9 @@ std::string GetSmName(se::CudaComputeCapability compute_capability) {
   // On Hopper, default to sm_90a so that all instructions can be used. But
   // only sm_90 is forward compatible, so don't use sm_90a with newer hardware:
   // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility
+  // Similarly for sm_100a (Blackwell).
   absl::string_view extension =
-      (compute_capability.major == 9 && sm_version == 90) ? "a" : "";
+      stream_executor::ShouldUsePtxExtension(compute_capability) ? "a" : "";
   return absl::StrCat("sm_", sm_version, extension);
 }
 
@@ -331,7 +333,7 @@ absl::StatusOr<std::string> CompileToPtx(
 
 namespace {
 constexpr stream_executor::SemanticVersion kFallbackPtxVersion{6, 5, 0};
-constexpr stream_executor::SemanticVersion kMaxPtxVersion{8, 5, 0};
+constexpr stream_executor::SemanticVersion kMaxPtxVersion{8, 6, 0};
 }  // namespace
 
 stream_executor::SemanticVersion
@@ -354,6 +356,10 @@ DetermineHighestSupportedPtxVersionFromCudaVersion(
   if (cuda_version < stream_executor::SemanticVersion{12, 6, 0}) {
     return {cuda_version.major() - 4, cuda_version.minor(), 0};
   }
+  // CUDA 12.6 -> PTX 8.5
+  if (cuda_version < stream_executor::SemanticVersion{12, 7, 0}) {
+    return {cuda_version.major() - 4, cuda_version.minor() - 1, 0};
+  }
 
   // Return maximum known PTX version.
   return kMaxPtxVersion;

xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc

@@ -27,13 +27,11 @@ namespace {
 namespace se = ::stream_executor;
 
 TEST(UtilsTest, TestGetSmName) {
-  se::CudaComputeCapability cc_hopper(9, 0);
-  ASSERT_EQ(nvptx::GetSmName(cc_hopper), "sm_90a");
-  // Do not default to sm90_a after Hopper, because it is not forward
-  // compatible.
-  // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility
-  se::CudaComputeCapability cc_next(10, 0);
-  ASSERT_EQ(nvptx::GetSmName(cc_next), "sm_90");
+  ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability{9, 0}), "sm_90a");
+  ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability{10, 0}), "sm_100a");
+  // Do not use the extension for a yet-unknown compute capability.
+  // https://docs.nvidia.com/cuda/parallel-thread-execution/#release-notes-ptx-release-history
+  ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability{10, 9}), "sm_100");
 }
 
 using VersionPair = std::pair<se::SemanticVersion, se::SemanticVersion>;

xla/stream_executor/cuda/driver_compilation_provider.cc

@@ -165,7 +165,7 @@ absl::StatusOr<Assembly> DriverCompilationProvider::CompileAndLink(
   CHECK(info_log_buffer_size() <= kInfoLogBufferSize);
   info_log_buffer.resize(info_log_buffer_size());
 
-  absl::string_view extension = (cc.major == 9 && cc.minor == 0) ? "a" : "";
+  absl::string_view extension = ShouldUsePtxExtension(cc) ? "a" : "";
   std::string architecture = absl::StrCat("sm_", cc.major, cc.minor, extension);
 
   if (result != CUDA_SUCCESS) {

xla/stream_executor/cuda/nvjitlink_impl.cc

@@ -80,7 +80,6 @@ static absl::Status ToStatus(nvJitLinkResult status,
     }                                                                    \
   } while (false)
 
-
 static absl::StatusOr<std::string> nvJitLinkGetErrorLog(
     nvJitLinkHandle link_handle) {
   size_t size{};
@@ -139,7 +138,7 @@ absl::StatusOr<std::vector<uint8_t>> CompileAndLinkUsingLibNvJitLink(
   // On Hopper, default to sm_90a so that all instructions can be used. But
   // only sm_90 is forward compatible, so don't use sm_90a with newer hardware:
   // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility
-  absl::string_view extension = (cc.major == 9 && cc.minor == 0) ? "a" : "";
+  absl::string_view extension = ShouldUsePtxExtension(cc) ? "a" : "";
   std::string architecture = absl::StrCat("sm_", cc.major, cc.minor, extension);
   cli_args.emplace_back(absl::StrCat("-arch=", architecture));
 

xla/stream_executor/cuda/ptx_compiler_helpers.cc

@@ -101,4 +101,10 @@ void WarnIfBadPtxasVersion(absl::string_view method,
   });
 }
 
+// The extension is used for compute capabilities 9.0 and 10.0.
+// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility
+bool ShouldUsePtxExtension(const CudaComputeCapability& cc) {
+  return (cc.major == 9 && cc.minor == 0) || (cc.major == 10 && cc.minor == 0);
+}
+
 }  // namespace stream_executor

xla/stream_executor/cuda/ptx_compiler_helpers.h

@@ -43,6 +43,15 @@ absl::Status CreateErrorFromPTXASLog(absl::string_view log,
 void WarnIfBadPtxasVersion(absl::string_view method,
                            const CudaComputeCapability& cc,
                            SemanticVersion compiler_version);
+
+// Determine whether the PTX extension for a compute capability should be used.
+//
+// Returns true if the argument compute capability has PTX extensions that are
+// only valid for that compute capability. For example, "sm_90" only includes
+// features that are forward compatible, whereas "sm_90a" (the extension) also
+// includes Hopper-specific features, such as WGMMA. We want to use the latter.
+bool ShouldUsePtxExtension(const CudaComputeCapability& cc);
+
 }  // namespace stream_executor
 
 #endif  // XLA_STREAM_EXECUTOR_CUDA_PTX_COMPILER_HELPERS_H_

xla/stream_executor/cuda/ptx_compiler_impl.cc

@@ -97,7 +97,7 @@ absl::StatusOr<std::vector<uint8_t>> CompileGpuAsmUsingLibNvPtxCompiler(
   // On Hopper, default to sm_90a so that all instructions can be used. But
   // only sm_90 is forward compatible, so don't use sm_90a with newer hardware:
   // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility
-  absl::string_view extension = (cc.major == 9 && cc.minor == 0) ? "a" : "";
+  absl::string_view extension = ShouldUsePtxExtension(cc) ? "a" : "";
   std::string architecture = absl::StrCat("sm_", cc.major, cc.minor, extension);
 
   options.extra_flags.emplace_back(absl::StrCat("-arch=", architecture));

xla/stream_executor/cuda/subprocess_compilation.cc

@@ -294,7 +294,7 @@ absl::StatusOr<std::vector<uint8_t>> CompileGpuAsmUsingPtxAs(
   // On Hopper, default to sm_90a so that all instructions can be used. But
   // only sm_90 is forward compatible, so don't use sm_90a with newer hardware:
   // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#ptx-compatibility
-  std::string extension = (cc.major == 9 && cc.minor == 0) ? "a" : "";
+  std::string extension = ShouldUsePtxExtension(cc) ? "a" : "";
   std::vector<std::string> ptxas_args = {
       std::string{ptxas_path},
       ptx_path,
@@ -515,7 +515,7 @@ absl::StatusOr<std::vector<uint8_t>> LinkUsingNvlink(
   };
   std::vector<std::string> args;
   args.push_back(std::string{nvlink_path});
-  absl::string_view extension = (cc.major == 9 && cc.minor == 0) ? "a" : "";
+  absl::string_view extension = ShouldUsePtxExtension(cc) ? "a" : "";
   args.push_back(absl::StrCat("-arch=sm_", cc.major, cc.minor, extension));
   for (int i = 0; i < images.size(); i++) {
     args.push_back(temp_files[i]);