Fix handling of the 9.0a and 12.0a capabilities (#202)

The code that computes the intersection between a kernel's capabilities
and the supported capabilities failed because e.g. the base capability
of 9.0a (9.0) did not match 9.0+PTX.

Instead we should just leave the supported capabilities at regular
version numbers. Instead, make a list of default kernel capabilities and
use that and intersect with the supported capabilities if no
capabilities are specified for a kernel.

Author: danieldk

build2cmake/src/cuda_supported_archs.json

@@ -1,13 +1 @@
-[
-  "7.0",
-  "7.2",
-  "7.5",
-  "8.0",
-  "8.6",
-  "8.7",
-  "8.9",
-  "9.0+PTX",
-  "10.0",
-  "10.1",
-  "12.0+PTX"
-]
+["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0", "10.0", "10.1", "12.0"]

build2cmake/src/templates/cuda/kernel.cmake

@@ -18,7 +18,7 @@ if(GPU_LANG STREQUAL "CUDA")
   {% if cuda_capabilities %}
     cuda_archs_loose_intersection({{kernel_name}}_ARCHS "{{ cuda_capabilities|join(";") }}" "${CUDA_ARCHS}")
   {% else %}
-    cuda_archs_loose_intersection({{kernel_name}}_ARCHS "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection({{kernel_name}}_ARCHS "${CUDA_DEFAULT_KERNEL_ARCHS}" "${CUDA_ARCHS}")
   {% endif %}
   message(STATUS "Capabilities for kernel {{kernel_name}}: {{ '${' + kernel_name + '_ARCHS}'}}")
   set_gencode_flags_for_srcs(SRCS {{'"${' + kernel_name + '_SRC}"'}} CUDA_ARCHS "{{ '${' + kernel_name + '_ARCHS}'}}")

build2cmake/src/templates/cuda/preamble.cmake

@@ -34,6 +34,13 @@ if (NOT TARGET_DEVICE STREQUAL "cuda" AND
     return()
 endif()
 
+if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
+   CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
+ set(CUDA_DEFAULT_KERNEL_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0+PTX")
+else()
+  set(CUDA_DEFAULT_KERNEL_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0+PTX")
+endif()
+
 if (NOT HIP_FOUND AND CUDA_FOUND)
   set(GPU_LANG "CUDA")
 

lib/torch-extension/default.nix

@@ -132,9 +132,9 @@ stdenv.mkDerivation (prevAttrs: {
       CUDAToolkit_ROOT = "${lib.getDev cudaPackages.cuda_nvcc}";
       TORCH_CUDA_ARCH_LIST =
         if cudaPackages.cudaOlder "12.8" then
-          "7.0;7.5;8.0;8.6;8.9;9.0+PTX"
+          "7.0;7.5;8.0;8.6;8.9;9.0"
         else
-          "7.0;7.5;8.0;8.6;8.9;9.0;10.0;10.1;12.0+PTX";
+          "7.0;7.5;8.0;8.6;8.9;9.0;10.0;10.1;12.0";
     }
     // lib.optionalAttrs rocmSupport {
       PYTORCH_ROCM_ARCH = lib.concatStringsSep ";" torch.rocmArchs;