KernelInterface: sub-group shuffles, votes and a constant width

.buildkite/pipeline.yml

@@ -14,19 +14,27 @@ steps:
     steps:
       # Julia 1.10 doesn't support [sources] in Project.toml,
       # so we need to manually dev workspace packages before instantiating.
-      - label: "Julia 1.10"
+      # KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet, so they need
+      # to be checked out from git and developed in the same Pkg operation, because
+      # CUDACore's compat excludes the registered versions.
+      # Julia 1.11 doesn't support workspaces, so it doesn't pick up KernelAbstractions
+      # from the test project's [sources] when instantiating, and needs the same treatment.
+      - label: "Julia {{matrix.julia}}"
         plugins:
           - JuliaCI/julia#v1:
-              version: "1.10"
+              version: "{{matrix.julia}}"
           - JuliaCI/julia-coverage#v1:
               dirs: ["."]
         agents:
           queue: "cuda"
         commands: |
           echo -e "[CUDA_Runtime_jll]\nlocal = \"true\"\n\n[CUDA_Compiler_jll]\nlocal = \"true\"" >LocalPreferences.toml
+          # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet
+          git clone --depth 1 --branch vc/ki-subgroup-ops https://github.com/JuliaGPU/KernelAbstractions.jl ka
           julia --project=. -e '
             using Pkg
             Pkg.develop([PackageSpec(; path) for path in (
+              "ka", "ka/lib/KernelInterface",
               "CUDACore", "CUDATools", "lib/cupti", "lib/nvml",
               "lib/cublas", "lib/cusparse", "lib/cusolver", "lib/cufft", "lib/curand",
               "lib/cudnn", "lib/cutensor", "lib/cutensornet", "lib/custatevec"
@@ -37,6 +45,11 @@ steps:
           (build.message =~ /\[only [^\]]*(tests|julia)/ ||
            build.message !~ /\[only / && !build.pull_request.draft)
         timeout_in_minutes: 90
+        matrix:
+          setup:
+            julia:
+              - "1.10"
+              - "1.11"
 
       # the newest release, also tested on draft pull requests
       - &julia
@@ -71,7 +84,6 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.11"
               - "1.12"
 
       - <<: *julia
@@ -150,7 +162,9 @@ steps:
       - label: "Enzyme.jl"
         plugins:
           - JuliaCI/julia#v1:
-              version: "1.11"
+              # XXX: needs Julia 1.12+ to pick up KernelAbstractions 0.10 from the test
+              #      project's [sources]
+              version: "1.12"
           - JuliaCI/julia-coverage#v1:
               dirs: ["."]
         commands: |
@@ -166,16 +180,13 @@ steps:
               # to check compatibility, also add Enzyme to the main environment
               # (or Pkg.test, which merges both environments, could fail)
               Pkg.activate(".")
-              # Try to co-develop Enzyme and KA, if that fails, try just to dev Enzyme
+              # XXX: KernelAbstractions 0.10 has no Enzyme integration to co-develop, and
+              #      comes from the test project's [sources] until it is registered
               try
-                Pkg.develop([PackageSpec("Enzyme"), PackageSpec("KernelAbstractions")])
+                Pkg.develop([PackageSpec("Enzyme")])
               catch err
-                try
-                  Pkg.develop([PackageSpec("Enzyme")])
-                catch err
-                  @error "Could not install Enzyme" exception=(err,catch_backtrace())
-                  exit(3)
-                end
+                @error "Could not install Enzyme" exception=(err,catch_backtrace())
+                exit(3)
               end
             end
 
@@ -196,7 +207,9 @@ steps:
       - label: "GPU-less environment"
         plugins:
           - JuliaCI/julia#v1:
-              version: "1.11"
+              # instantiating the workspace needs Julia 1.12+ to pick up
+              # KernelAbstractions 0.10 from the test project's [sources]
+              version: "1.12"
           - JuliaCI/julia-coverage#v1:
               dirs: ["."]
           - JuliaCI/julia-test#v1:

CUDACore/Project.toml

@@ -14,7 +14,7 @@ ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LLVMLoopInfo = "8b046642-f1f6-4319-8d3c-209ddc03c586"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
@@ -29,19 +29,25 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
 ChainRulesCoreExt = "ChainRulesCore"
 EnzymeCoreExt = "EnzymeCore"
+KernelAbstractionsExt = "KernelAbstractions"
 SpecialFunctionsExt = "SpecialFunctions"
+StaticArraysExt = "StaticArrays"
+
+[sources]
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"}
 
 [compat]
 Adapt = "4.4"
@@ -58,7 +64,8 @@ ExprTools = "0.1"
 GPUArrays = "12.1"
 GPUCompiler = "2.13"
 GPUToolbox = "3.3.1"
-KernelAbstractions = "0.9.44"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "10"
 LLVMLoopInfo = "1"
 LazyArtifacts = "1"
@@ -82,3 +89,4 @@ julia = "1.10"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

CUDACore/ext/KernelAbstractionsExt.jl

@@ -0,0 +1,29 @@
+module KernelAbstractionsExt
+
+using CUDACore
+using CUDACore: GPUArrays
+
+import KernelAbstractions as KA
+
+import Adapt
+
+# TODO: Move AbstractGPUSparseArray stuff out
+Adapt.adapt_storage(::KA.CPU, a::Union{CuArray,GPUArrays.AbstractGPUSparseArray}) = Adapt.adapt(Array, a)
+
+## kernel launch
+
+# KernelAbstractions launches kernels through the KernelInterface back-end; tell the
+# compiler about statically sized workgroups
+function KA.compiler_options(obj::KA.Kernel{CUDABackend})
+    if KA.workgroupsize(obj) <: KA.StaticSize
+        return (; maxthreads = prod(KA.get(KA.workgroupsize(obj))))
+    else
+        return (;)
+    end
+end
+
+## other
+
+Adapt.adapt_storage(to::KA.ConstAdaptor, a::CuDeviceArray) = Base.Experimental.Const(a)
+
+end

CUDACore/ext/StaticArraysExt.jl

@@ -2,8 +2,8 @@
 
 module StaticArraysExt
 
-using ..CUDACore
-using ..CUDACore: @device_override, @gputhrow
+using CUDACore
+using CUDACore: @device_override, @gputhrow
 
 import StaticArrays
 

CUDACore/src/CUDACore.jl

@@ -17,7 +17,8 @@ using LLVM, LLVM.IR, LLVM.Build
 using LLVM.Interop
 using Core: LLVMPtr
 
-import KernelAbstractions
+import KernelInterface
+import KernelInterface as KI
 
 using Adapt: Adapt, adapt, WrappedArray
 
@@ -127,15 +128,11 @@ include("iterator.jl")
 include("complex.jl")
 include("library_types.jl")
 
-# KernelAbstractions
+# KernelInterface
 include("CUDAKernels.jl")
 import .CUDAKernels: CUDABackend
 export CUDABackend
 
-# StaticArrays is still a direct dependency, so directly include the extension
-include("../ext/StaticArraysExt.jl")
-# NOTE: StaticArrays is a direct dep, so extension is directly included
-
 include("precompile.jl")
 
 end