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

.buildkite/pipeline.yml

@@ -70,25 +70,63 @@ steps:
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"
       JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
+  # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet. Without
+  #      workspaces (Julia 1.11), only the test project, which depends on KernelAbstractions,
+  #      can take it from [sources], and Julia 1.10 doesn't support [sources] at all. So run
+  #      the tests in the test project, developing the packages from git on Julia 1.10, in
+  #      one Pkg operation because AMDGPU's compat excludes the registered versions.
+  #
   # the oldest supported release
-  - <<: *julia
+  - &julia_test_project
+    label: "Julia {{matrix.version}}"
     matrix:
       setup:
         version:
           - "1.10"
-        test_args:
-          - "--verbose"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.version}}"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "rocm"
+      rocmgpu: "*"
     if: |
       build.message !~ /\[skip [^\]]*(tests|julia)/ &&
       (build.message =~ /\[only [^\]]*(tests|julia)/ ||
        build.message !~ /\[only / && !build.pull_request.draft)
+    commands: |
+      git clone --depth 1 --branch vc/ki-subgroup-ops https://github.com/JuliaGPU/KernelAbstractions.jl ka
+      julia --project=test -e '
+        using Pkg
+        if VERSION < v"1.11"
+          Pkg.develop([PackageSpec(; path) for path in (".", "ka", "ka/lib/KernelInterface")])
+        else
+          Pkg.instantiate()
+        end'
+      julia --project=test --check-bounds=yes --code-coverage=@. test/runtests.jl --verbose
+    timeout_in_minutes: 90
+    env:
+      JULIA_NUM_THREADS: 4
+      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
   # intermediate releases, not tested on pull requests
-  - <<: *julia
+  - <<: *julia_test_project
     matrix:
       setup:
         version:
           - "1.11"
+    if: |
+      build.message !~ /\[skip [^\]]*(tests|julia)/ &&
+      (build.message =~ /\[only [^\]]*(tests|julia)/ ||
+       build.message !~ /\[only / && build.pull_request.id == null)
+
+  - <<: *julia
+    matrix:
+      setup:
+        version:
           - "1.12"
         test_args:
           - "--verbose"
@@ -145,7 +183,8 @@ steps:
   - label: "GPU-less environment"
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.10"
+          # XXX: needs Julia 1.12+ to pick up KernelAbstractions 0.10 from [sources]
+          version: "1.12"
       - JuliaCI/julia-test#v1:
           run_tests: false
     command: |

Project.toml

@@ -17,7 +17,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"
 LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
@@ -31,22 +31,27 @@ Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
-StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
 AMDGPUChainRulesCoreExt = "ChainRulesCore"
 AMDGPUEnzymeCoreExt = "EnzymeCore"
+AMDGPUKernelAbstractionsExt = "KernelAbstractions"
 AMDGPUSparseMatricesCSRExt = "SparseMatricesCSR"
 AMDGPUSpecialFunctionsExt = "SpecialFunctions"
 
+[sources]
+KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"}
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"}
+
 [compat]
 AMDGPU_LLVM_Backend_jll = "23"
 AbstractFFTs = "1.0"
@@ -60,7 +65,8 @@ ExprTools = "0.1"
 GPUArrays = "12.1"
 GPUCompiler = "2.13"
 GPUToolbox = "3.3.1"
-KernelAbstractions = "0.9.2"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "10"
 LLVMDowngrader_jll = "0.12"
 Libdl = "1"
@@ -76,7 +82,9 @@ Scratch = "1"
 SparseArrays = "1"
 SparseMatricesCSR = "0.6.9"
 SpecialFunctions = "2"
-StaticArraysCore = "1"
 Statistics = "1"
 UnsafeAtomics = "0.4"
 julia = "1.10"
+
+[extras]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"

docs/src/api/devices.md

@@ -56,6 +56,7 @@ AMDGPU.device_id!
 ```@docs
 AMDGPU.HIP.name
 AMDGPU.HIP.wavefrontsize
+AMDGPU.HIP.max_workgroup_dims
 AMDGPU.HIP.gcn_arch
 AMDGPU.HIP.device_id
 AMDGPU.HIP.properties

docs/src/api/intrinsics.md

@@ -23,6 +23,7 @@ AMDGPU.Device.blockDim
 
 ```@docs
 AMDGPU.sync_workgroup
+AMDGPU.sync_wavefront
 AMDGPU.sync_workgroup_count
 AMDGPU.sync_workgroup_and
 AMDGPU.sync_workgroup_or

ext/AMDGPUKernelAbstractionsExt.jl

@@ -0,0 +1,33 @@
+module AMDGPUKernelAbstractionsExt
+
+import AMDGPU
+using AMDGPU: GPUArrays, ROCBackend
+
+import Adapt
+import KernelAbstractions as KA
+import LLVM
+
+Adapt.adapt_storage(::KA.CPU, a::Union{AMDGPU.ROCArray, 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{ROCBackend})
+    if KA.workgroupsize(obj) <: KA.StaticSize
+        return (; maxthreads = prod(KA.get(KA.workgroupsize(obj))))
+    else
+        return (;)
+    end
+end
+
+## other
+
+# `@Const` arrays are read through the constant address space
+function Adapt.adapt_storage(::KA.ConstAdaptor, a::AMDGPU.ROCDeviceArray{T}) where T
+    ptr = LLVM.Interop.addrspacecast(Core.LLVMPtr{T,AMDGPU.Device.AS.Constant}, a.ptr)
+    AMDGPU.ROCDeviceArray(a.dims, ptr)
+end
+
+end

src/AMDGPU.jl

@@ -11,6 +11,7 @@ using LLVM
 using Preferences
 using Printf
 
+import KernelInterface
 import UnsafeAtomics
 import Atomix
 import Atomix: @atomic, @atomicswap, @atomicreplace
@@ -79,9 +80,11 @@ include("memory.jl")
 
 Base.Experimental.@MethodTable(method_table)
 
-# Needs to be before Device since sync uses this. GPUCompiler emits `device` as LLVM's `agent`.
+# Needs to be before Device since sync uses this. GPUCompiler emits `device` as LLVM's
+# `agent`, and `subgroup` as `wavefront`.
 const syncscope_agent = UnsafeAtomics.device
 const syncscope_workgroup = UnsafeAtomics.workgroup
+const syncscope_wavefront = UnsafeAtomics.subgroup
 
 # Referenced by the generated `kernel_state()`, and generators run in the world they are
 # defined in, so this has to precede the device code.
@@ -102,7 +105,7 @@ import .Device: ROCDeviceArray, AS, HostCall, HostCallHolder, hostcall!
 import .Device: @ROCDynamicLocalArray, @ROCStaticLocalArray
 import .Device: workitemIdx, workgroupIdx, workgroupDim, gridItemDim, gridGroupDim
 import .Device: threadIdx, blockIdx, blockDim
-import .Device: sync_workgroup, sync_workgroup_count, sync_workgroup_and, sync_workgroup_or
+import .Device: sync_workgroup, sync_wavefront, sync_workgroup_count, sync_workgroup_and, sync_workgroup_or
 import .Device: @rocprint, @rocprintln, @rocprintf
 
 export ROCDeviceArray, @ROCDynamicLocalArray, @ROCStaticLocalArray
@@ -132,6 +135,7 @@ include("dnn/MIOpen.jl")
 
 include("random.jl")
 
+# KernelInterface
 include("ROCKernels.jl")
 import .ROCKernels: ROCBackend
 export ROCBackend