Add support for thread block clusters

.buildkite/pipeline.yml

@@ -30,10 +30,11 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.10"
-              - "1.11"
-              - "1.12"
-              - "nightly"
+              #TODO - "1.10"
+              #TODO - "1.11"
+              #TODO - "1.12"
+              - "1.13"
+              #TODO - "nightly"
           adjustments:
             - with:
                 julia: "nightly"
@@ -68,16 +69,17 @@ steps:
         matrix:
           setup:
             cuda:
-              - "13.0"
-              - "12.9"
-              - "12.8"
-              - "12.6"
-              - "12.5"
-              - "12.4"
-              - "12.3"
-              - "12.2"
-              - "12.1"
-              - "12.0"
+              - "13.1"
+              #TODO - "13.0"
+              #TODO - "12.9"
+              #TODO - "12.8"
+              #TODO - "12.6"
+              #TODO - "12.5"
+              #TODO - "12.4"
+              #TODO - "12.3"
+              #TODO - "12.2"
+              #TODO - "12.1"
+              #TODO - "12.0"
           adjustments:
             - with:
                 cuda: "13.0"

lib/cudadrv/execution.jl

@@ -44,22 +44,26 @@ end
 
 """
     launch(f::CuFunction; args...; blocks::CuDim=1, threads::CuDim=1,
-           cooperative=false, shmem=0, stream=stream())
+           clustersize::CuDim=1, cooperative=false, shmem=0, stream=stream())
 
 Low-level call to launch a CUDA function `f` on the GPU, using `blocks` and `threads` as
 respectively the grid and block configuration. Dynamic shared memory is allocated according
-to `shmem`, and the kernel is launched on stream `stream`.
+to `shmem`, and the kernel is launched on stream `stream`. If `clustersize > 1` and compute
+capability is `>= 9.0`, [thread block clusters](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#thread-block-clusters)
+are launched. If `clustersize > 1` and compute capability is `< 9.0`, an error is thrown, as
+thread block clusters are not supported.
 
 Arguments to a kernel should either be bitstype, in which case they will be copied to the
 internal kernel parameter buffer, or a pointer to device memory.
 
 This is a low-level call, prefer to use [`cudacall`](@ref) instead.
 """
 function launch(f::CuFunction, args::Vararg{Any,N}; blocks::CuDim=1, threads::CuDim=1,
-                cooperative::Bool=false, shmem::Integer=0,
+                clustersize::CuDim=1, cooperative::Bool=false, shmem::Integer=0,
                 stream::CuStream=stream()) where {N}
     blockdim = CuDim3(blocks)
     threaddim = CuDim3(threads)
+    clusterdim = CuDim3(clustersize)
 
     try
         pack_arguments(args...) do kernelParams
@@ -68,29 +72,46 @@ function launch(f::CuFunction, args::Vararg{Any,N}; blocks::CuDim=1, threads::Cu
                                           blockdim.x, blockdim.y, blockdim.z,
                                           threaddim.x, threaddim.y, threaddim.z,
                                           shmem, stream, kernelParams)
-            else
+            elseif clusterdim.x == 1 && clusterdim.y == 1 && clusterdim.z == 1
                 cuLaunchKernel(f,
                                blockdim.x, blockdim.y, blockdim.z,
                                threaddim.x, threaddim.y, threaddim.z,
                                shmem, stream, kernelParams, C_NULL)
+            else
+                attr_ref = Ref{CUDA.CUlaunchAttribute}()
+                GC.@preserve attr_ref stream begin
+                    attr = Base.unsafe_convert(Ptr{CUDA.CUlaunchAttribute}, attr_ref)
+                    attr.id = CUDA.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
+                    attr.value.clusterDim.x = clusterdim.x
+                    attr.value.clusterDim.y = clusterdim.y
+                    attr.value.clusterDim.z = clusterdim.z
+                    config = Ref(CUlaunchConfig(blockdim.x, blockdim.y, blockdim.z,
+                                                threaddim.x, threaddim.y, threaddim.z,
+                                                shmem, stream.handle, attr, 1))
+                    cuLaunchKernelEx(config, f, kernelParams, C_NULL)
+                end
             end
         end
     catch err
-        diagnose_launch_failure(f, err; blockdim, threaddim, shmem)
+        diagnose_launch_failure(f, err; blockdim, threaddim, clusterdim, shmem)
     end
 end
 
-@noinline function diagnose_launch_failure(f::CuFunction, err; blockdim, threaddim, shmem)
+@noinline function diagnose_launch_failure(f::CuFunction, err; blockdim, threaddim, clusterdim, shmem)
     if !isa(err, CuError) || !in(err.code, [ERROR_INVALID_VALUE,
                                             ERROR_LAUNCH_OUT_OF_RESOURCES])
         rethrow()
     end
 
     # essentials
     (blockdim.x>0 && blockdim.y>0 && blockdim.z>0) ||
-        error("Grid dimensions should be non-null")
+        error("Grid dimensions $blockdim are not positive")
     (threaddim.x>0 && threaddim.y>0 && threaddim.z>0) ||
-        error("Block dimensions should be non-null")
+        error("Block dimensions $threaddim are not positive")
+    (clusterdim.x>0 && clusterdim.y>0 && clusterdim.z>0) ||
+        error("Cluster dimensions $clusterdim are not positive")
+    (blockdim.x % clusterdim.x == 0 && blockdim.y % clusterdim.y == 0 && blockdim.z % clusterdim.z == 0) ||
+        error("Block dimensions $blockdim are not multiples of the cluster dimensions $clusterdim")
 
     # check device limits
     dev = device()

lib/nvml/NVML.jl

@@ -12,22 +12,23 @@ import Libdl
 
 export has_nvml
 
-function libnvml()
-    @memoize begin
-        if Sys.iswindows()
-            # the NVSMI dir isn't added to PATH by the installer
-            nvsmi = joinpath(ENV["ProgramFiles"], "NVIDIA Corporation", "NVSMI")
-            if isdir(nvsmi)
-                joinpath(nvsmi, "nvml.dll")
-            else
-                # let's just hope for the best
-                "nvml"
-            end
-        else
-            "libnvidia-ml.so.1"
-        end
-    end::String
-end
+#TODO function libnvml()
+#TODO     @memoize begin
+#TODO         if Sys.iswindows()
+#TODO             # the NVSMI dir isn't added to PATH by the installer
+#TODO             nvsmi = joinpath(ENV["ProgramFiles"], "NVIDIA Corporation", "NVSMI")
+#TODO             if isdir(nvsmi)
+#TODO                 joinpath(nvsmi, "nvml.dll")
+#TODO             else
+#TODO                 # let's just hope for the best
+#TODO                 "nvml"
+#TODO             end
+#TODO         else
+#TODO             "libnvidia-ml.so.1"
+#TODO         end
+#TODO     end::String
+#TODO end
+const libnvml::String = "libnvidia-ml.so.1"
 
 function has_nvml()
     @memoize begin
@@ -37,7 +38,7 @@ function has_nvml()
             return false
         end
 
-        if Libdl.dlopen(libnvml(); throw_error=false) === nothing
+        if Libdl.dlopen(libnvml; throw_error=false) === nothing
             return false
         end