Merge branch 'main' into vc/printf

Author: vchuravy

.buildkite/pipeline.yml

@@ -11,10 +11,15 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |
+      julia -e 'println("--- :julia: Developing CUDA")
+                using Pkg
+                Pkg.develop(; name="CUDA")
+                Pkg.develop(; name="GPUArrays")'
+      sed -i 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/CUDA/Project.toml
+      sed -i 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/GPUArrays/Project.toml
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
-                Pkg.develop(; path=pwd())
-                Pkg.develop(; name="CUDA")' || exit 3
+                Pkg.develop(; path=pwd())' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
@@ -25,38 +30,40 @@ steps:
     timeout_in_minutes: 120
     soft_fail:
       - exit_status: 3
+    env:
+      JULIA_PROJECT: "@cuda"
 
-  - label: "CUDA Enzyme Julia {{matrix.version}}"
-    matrix:
-      setup:
-        version:
-          - "1.10"
-          - "1.11"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "{{matrix.version}}"
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
-    command: |
-      julia -e 'println("--- :julia: Instantiating project")
-                using Pkg
-                try
-                   Pkg.develop([PackageSpec(; path=pwd()), PackageSpec("Enzyme"), PackageSpec("EnzymeCore"), PackageSpec("CUDA")])
-                catch err
-                   Pkg.develop(; path=pwd())
-                  Pkg.add(["CUDA", "Enzyme"])
-                end' || exit 3
+  # - label: "CUDA Enzyme Julia {{matrix.version}}"
+  #   matrix:
+  #     setup:
+  #       version:
+  #         - "1.10"
+  #         - "1.11"
+  #   plugins:
+  #     - JuliaCI/julia#v1:
+  #         version: "{{matrix.version}}"
+  #     - JuliaCI/julia-coverage#v1:
+  #         codecov: true
+  #   command: |
+  #     julia -e 'println("--- :julia: Instantiating project")
+  #               using Pkg
+  #               try
+  #                  Pkg.develop([PackageSpec(; path=pwd()), PackageSpec("Enzyme"), PackageSpec("EnzymeCore"), PackageSpec("CUDA")])
+  #               catch err
+  #                  Pkg.develop(; path=pwd())
+  #                 Pkg.add(["CUDA", "Enzyme"])
+  #               end' || exit 3
 
-      julia -e 'println("+++ :julia: Running tests")
-                using CUDA
-                include("test/extensions/enzyme.jl")
-                enzyme_testsuite(CUDABackend, CuArray, true)'
-    agents:
-      queue: "juliagpu"
-      cuda: "*"
-    timeout_in_minutes: 120
-    soft_fail:
-      - exit_status: 3
+  #     julia -e 'println("+++ :julia: Running tests")
+  #               using CUDA
+  #               include("test/extensions/enzyme.jl")
+  #               enzyme_testsuite(CUDABackend, CuArray, true)'
+  #   agents:
+  #     queue: "juliagpu"
+  #     cuda: "*"
+  #   timeout_in_minutes: 120
+  #   soft_fail:
+  #     - exit_status: 3
 
   - label: "Metal Julia {{matrix.version}}"
     matrix:
@@ -70,10 +77,15 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |
+      julia -e 'println("--- :julia: Developing Metal")
+                using Pkg
+                Pkg.develop(; name="Metal")
+                Pkg.develop(; name="GPUArrays")'
+      sed -i '' 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/Metal/Project.toml
+      sed -i '' 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/GPUArrays/Project.toml
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
-                Pkg.develop(; path=pwd())
-                Pkg.develop(; name="Metal")' || exit 3
+                Pkg.develop(; path=pwd())' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
@@ -85,6 +97,8 @@ steps:
     timeout_in_minutes: 120
     soft_fail:
       - exit_status: 3
+    env:
+      JULIA_PROJECT: "@metal"
 
   - label: "oneAPI Julia {{matrix.version}}"
     matrix:
@@ -98,10 +112,15 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |
+      julia -e 'println("--- :julia: Developing oneAPI")
+                using Pkg
+                Pkg.develop(; name="oneAPI")
+                Pkg.develop(; name="GPUArrays")'
+      sed -i 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/oneAPI/Project.toml
+      sed -i 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/GPUArrays/Project.toml
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
-                Pkg.develop(; path=pwd())
-                Pkg.develop(; name="oneAPI")' || exit 3
+                Pkg.develop(; path=pwd())' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
@@ -112,6 +131,8 @@ steps:
     timeout_in_minutes: 120
     soft_fail:
       - exit_status: 3
+    env:
+      JULIA_PROJECT: "@oneAPI"
 
   - label: "AMDGPU Julia {{matrix.version}}"
     matrix:
@@ -125,10 +146,15 @@ steps:
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: |
+      julia -e 'println("--- :julia: Developing AMDGPU")
+                using Pkg
+                Pkg.develop(; name="AMDGPU")
+                Pkg.develop(; name="GPUArrays")'
+      sed -i 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/AMDGPU/Project.toml
+      sed -i 's/^KernelAbstractions = "0\.9.*"/KernelAbstractions = "0.10"/' \${JULIA_DEPOT_PATH}/dev/GPUArrays/Project.toml
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
-                Pkg.develop(; path=pwd())
-                Pkg.develop(; name="AMDGPU")' || exit 3
+                Pkg.develop(; path=pwd())' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
@@ -140,6 +166,7 @@ steps:
     soft_fail:
       - exit_status: 3
     env:
+      JULIA_PROJECT: "@amdgpu"
       JULIA_NUM_THREADS: 4
 
 env:

Project.toml

@@ -14,6 +14,8 @@ OpenCL_jll = "6cb37087-e8b6-5417-8430-1f242f1e46e4"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SPIRVIntrinsics = "71d1d633-e7e8-4a92-83a1-de8814b09ba8"
+SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c"
+SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd"
@@ -32,13 +34,15 @@ SparseArraysExt = "SparseArrays"
 Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "0.1, 1"
 EnzymeCore = "0.7, 0.8.1"
-GPUCompiler = "1"
+GPUCompiler = "1.2"
 InteractiveUtils = "1.6"
 LLVM = "9"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"
 PrecompileTools = "1"
-SPIRVIntrinsics = "0.2.1"
+SPIRVIntrinsics = "0.3"
+SPIRV_LLVM_Backend_jll = "20"
+SPIRV_Tools_jll = "2024.4, 2025.1"
 SparseArrays = "<0.0.1, 1.6"
 StaticArrays = "0.12, 1.0"
 UUIDs = "<0.0.1, 1.6"

benchmark/benchmarks.jl

@@ -9,13 +9,16 @@ using Random
 
 if !haskey(ENV, "KA_BACKEND")
     const BACKEND = CPU()
+    const Ts = (Float32, Float64)
 else
     backend = ENV["KA_BACKEND"]
     if backend == "CPU"
         const BACKEND = CPU()
+        const Ts = (Float32, Float64)
     elseif backend == "CUDA"
         using CUDA
         const BACKEND = CUDABackend()
+        const Ts = (Float16, Float32, Float64)
     else
         error("Backend $backend not recognized")
     end
@@ -31,7 +34,7 @@ end
 SUITE["saxpy"] = BenchmarkGroup()
 
 let static = BenchmarkGroup()
-    for T in (Float16, Float32, Float64)
+    for T in Ts
         dtype = BenchmarkGroup()
         for N in (64, 256, 512, 1024, 2048, 4096, 16384, 32768, 65536, 262144, 1048576)
             dtype[N] = @benchmarkable begin
@@ -49,7 +52,7 @@ let static = BenchmarkGroup()
 end
 
 let default = BenchmarkGroup()
-    for T in (Float16, Float32, Float64)
+    for T in Ts
         dtype = BenchmarkGroup()
         for N in (64, 256, 512, 1024, 2048, 4096, 16384, 32768, 65536, 262144, 1048576)
             dtype[N] = @benchmarkable begin

examples/histogram.jl

@@ -5,31 +5,29 @@ include(joinpath(dirname(pathof(KernelAbstractions)), "../examples/utils.jl")) #
 
 # Function to use as a baseline for CPU metrics
 function create_histogram(input)
-    histogram_output = zeros(Int, maximum(input))
+    histogram_output = zeros(eltype(input), maximum(input))
     for i in input
         histogram_output[i] += 1
     end
     return histogram_output
 end
 
 # This a 1D histogram kernel where the histogramming happens on shmem
-@kernel function histogram_kernel!(histogram_output, input)
-    tid = @index(Global, Linear)
+@kernel unsafe_indices = true function histogram_kernel!(histogram_output, input)
+    gid = @index(Group, Linear)
     lid = @index(Local, Linear)
 
-    @uniform warpsize = Int(32)
-
-    @uniform gs = @groupsize()[1]
+    @uniform gs = prod(@groupsize())
+    tid = (gid - 1) * gs + lid
     @uniform N = length(histogram_output)
 
-    shared_histogram = @localmem Int (gs)
+    shared_histogram = @localmem eltype(input) (gs)
 
     # This will go through all input elements and assign them to a location in
     # shmem. Note that if there is not enough shem, we create different shmem
     # blocks to write to. For example, if shmem is of size 256, but it's
     # possible to get a value of 312, then we will have 2 separate shmem blocks,
     # one from 1->256, and another from 256->512
-    @uniform max_element = 1
     for min_element in 1:gs:N
 
         # Setting shared_histogram to 0
@@ -42,7 +40,7 @@ end
         end
 
         # Defining bin on shared memory and writing to it if possible
-        bin = input[tid]
+        bin = tid <= length(input) ? input[tid] : 0
         if bin >= min_element && bin < max_element
             bin -= min_element - 1
             @atomic shared_histogram[bin] += 1
@@ -58,10 +56,10 @@ end
 
 end
 
-function histogram!(histogram_output, input)
+function histogram!(histogram_output, input, groupsize = 256)
     backend = get_backend(histogram_output)
     # Need static block size
-    kernel! = histogram_kernel!(backend, (256,))
+    kernel! = histogram_kernel!(backend, (groupsize,))
     kernel!(histogram_output, input, ndrange = size(input))
     return
 end
@@ -74,9 +72,10 @@ function move(backend, input)
 end
 
 @testset "histogram tests" begin
-    rand_input = [rand(1:128) for i in 1:1000]
-    linear_input = [i for i in 1:1024]
-    all_two = [2 for i in 1:512]
+    # Use Int32 as some backends don't support 64-bit atomics
+    rand_input = Int32.(rand(1:128, 1000))
+    linear_input = Int32.(1:1024)
+    all_two = fill(Int32(2), 512)
 
     histogram_rand_baseline = create_histogram(rand_input)
     histogram_linear_baseline = create_histogram(linear_input)
@@ -86,14 +85,14 @@ end
     linear_input = move(backend, linear_input)
     all_two = move(backend, all_two)
 
-    rand_histogram = KernelAbstractions.zeros(backend, Int, 128)
-    linear_histogram = KernelAbstractions.zeros(backend, Int, 1024)
-    two_histogram = KernelAbstractions.zeros(backend, Int, 2)
+    rand_histogram = KernelAbstractions.zeros(backend, eltype(rand_input), Int(maximum(rand_input)))
+    linear_histogram = KernelAbstractions.zeros(backend, eltype(linear_input), Int(maximum(linear_input)))
+    two_histogram = KernelAbstractions.zeros(backend, eltype(all_two), Int(maximum(all_two)))
 
-    histogram!(rand_histogram, rand_input)
+    histogram!(rand_histogram, rand_input, 6)
     histogram!(linear_histogram, linear_input)
     histogram!(two_histogram, all_two)
-    KernelAbstractions.synchronize(CPU())
+    KernelAbstractions.synchronize(backend)
 
     @test isapprox(Array(rand_histogram), histogram_rand_baseline)
     @test isapprox(Array(linear_histogram), histogram_linear_baseline)

examples/memcopy.jl

@@ -16,8 +16,8 @@ function mycopy!(A, B)
     return
 end
 
-A = KernelAbstractions.zeros(backend, Float64, 128, 128)
-B = KernelAbstractions.ones(backend, Float64, 128, 128)
+A = KernelAbstractions.zeros(backend, f_type, 128, 128)
+B = KernelAbstractions.ones(backend, f_type, 128, 128)
 mycopy!(A, B)
 KernelAbstractions.synchronize(backend)
 @test A == B

examples/memcopy_static.jl

@@ -16,8 +16,8 @@ function mycopy_static!(A, B)
     return
 end
 
-A = KernelAbstractions.zeros(backend, Float64, 128, 128)
-B = KernelAbstractions.ones(backend, Float64, 128, 128)
+A = KernelAbstractions.zeros(backend, f_type, 128, 128)
+B = KernelAbstractions.ones(backend, f_type, 128, 128)
 mycopy_static!(A, B)
 KernelAbstractions.synchronize(backend)
 @test A == B

examples/performance.jl

@@ -26,7 +26,7 @@ end
 
 # Local memory variants
 
-@kernel function lmem_copy_kernel!(
+@kernel unsafe_indices = true function lmem_copy_kernel!(
         output, @Const(input),
         ::Val{BANK} = Val(1),
     ) where {BANK}
@@ -46,7 +46,7 @@ end
     @inbounds output[I, J] = tile[i, j]
 end
 
-@kernel function lmem_transpose_kernel!(
+@kernel unsafe_indices = true function lmem_transpose_kernel!(
         output, @Const(input),
         ::Val{BANK} = Val(1),
     ) where {BANK}
@@ -77,7 +77,7 @@ end
 
 # Local Memory + process multiple elements per lane
 
-@kernel function coalesced_copy_kernel!(
+@kernel unsafe_indices = true function coalesced_copy_kernel!(
         output, @Const(input),
         ::Val{BANK} = Val(1),
     ) where {BANK}
@@ -105,7 +105,7 @@ end
     end
 end
 
-@kernel function coalesced_transpose_kernel!(
+@kernel unsafe_indices = true function coalesced_transpose_kernel!(
         output, @Const(input),
         ::Val{BANK} = Val(1),
     ) where {BANK}

examples/performant_matmul.jl

@@ -4,9 +4,12 @@ using Test
 using Random
 include(joinpath(dirname(pathof(KernelAbstractions)), "../examples/utils.jl")) # Load backend
 
-const TILE_DIM = 32
+# We use a TILE_DIM of 16 as a safe value since while
+#  most backends support up to 1024 threads per group,
+#  Metal sometimes supports fewer.
+const TILE_DIM = 16
 
-@kernel function coalesced_matmul_kernel!(
+@kernel unsafe_indices = true function coalesced_matmul_kernel!(
         output, @Const(input1), @Const(input2), N, R, M,
         ::Val{BANK} = Val(1),
     ) where {BANK}