diff --git a/Project.toml b/Project.toml index 9ecaec47..8748f219 100644 --- a/Project.toml +++ b/Project.toml @@ -10,6 +10,7 @@ AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LLVM = "929cbde3-209d-540e-8aea-75f648917ca0" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -31,7 +32,8 @@ AcceleratedKernels = "0.5.8" Adapt = "4.7.3" GPUArraysCore = "= 0.2.2" JLD2 = "0.4, 0.5, 0.6" -KernelAbstractions = "0.9.43, 0.10" +KernelAbstractions = "0.10" +KernelInterface = "0.4.1" LLVM = "3.9, 4, 5, 6, 7, 8, 9, 10" LinearAlgebra = "1" Printf = "1" diff --git a/docs/src/interface.md b/docs/src/interface.md index 37663494..b45ae976 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -21,12 +21,12 @@ end ``` This will allow your defined type (in this case `JLArray`) to use the GPUArrays interface where available. -To be able to actually use the functionality that is defined for `AbstractGPUArray`s, you need to define the backend, like so: +To be able to actually use the functionality that is defined for `AbstractGPUArray`s, you need to define a [KernelInterface](https://github.com/JuliaGPU/KernelAbstractions.jl/tree/main/lib/KernelInterface) back-end, which compiles and launches kernels, like so: ```julia -import KernelAbstractions: Backend -struct CustomBackend <: KernelAbstractions.GPU -KernelAbstractions.get_backend(a::CA) where CA <: CustomArray = CustomBackend() +import KernelInterface as KI +struct CustomBackend <: KI.Backend end +KI.get_backend(a::CA) where CA <: CustomArray = CustomBackend() ``` There are numerous examples of potential interfaces for GPUArrays, such as with [JLArrays](https://github.com/JuliaGPU/GPUArrays.jl/blob/main/lib/JLArrays/src/JLArrays.jl), [CuArrays](https://github.com/JuliaGPU/CUDA.jl/blob/main/CUDACore/src/array.jl), and [ROCArrays](https://github.com/JuliaGPU/AMDGPU.jl/blob/main/src/array.jl). diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index d12d6762..3db7a050 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -16,6 +16,7 @@ using Reexport @reexport using GPUArraysCore using KernelAbstractions +import KernelInterface as KI import AcceleratedKernels as AK diff --git a/src/host/gemm.jl b/src/host/gemm.jl index d018806d..ee45f29b 100644 --- a/src/host/gemm.jl +++ b/src/host/gemm.jl @@ -33,42 +33,33 @@ end end const TILE_DIM = 16 -@kernel unsafe_indices = true function coalesced_matmul_kernel!( +function coalesced_matmul_kernel!( C, A, B, add, M, N, K, ::Val{TA}, ::Val{TB}, ::Val{TILE}, ) where {TA, TB, TILE} - @uniform TAT = eltype(A); @uniform TBT = eltype(B); @uniform R = eltype(C) + TAT = eltype(A); TBT = eltype(B); R = eltype(C) - # li, lj, _ = KI.get_local_id() - li, lj = @index(Local, NTuple) - - # grow, gcol, _ = KI.get_group_id() - grow, gcol = @index(Group, NTuple) - # Uncomment once moved to KernelInterface - # gi = (grow - 1) * TILE + li - # gj = (gcol - 1) * TILE + lj + li, lj, _ = KI.get_local_id() + grow, gcol, _ = KI.get_group_id() + gi = (grow - 1) * TILE + li + gj = (gcol - 1) * TILE + lj - # padded by one row to avoid bank conflicts on shared memory - tile1 = @localmem TAT (TILE + 1, TILE) - tile2 = @localmem TBT (TILE + 1, TILE) + # no padding: consecutive work-items (`li`) access consecutive elements of a column, or + # all the same element of `tile2`, so there are no bank conflicts to avoid. unpadded + # columns stay aligned, which lets each work-item read `tile2[:, lj]` with vector loads. + tile1 = KI.localmemory(TAT, (TILE, TILE)) + tile2 = KI.localmemory(TBT, (TILE, TILE)) # private variable for tile C - @uniform Tacc = promote_type(R, typeof(zero(TAT) * zero(TBT) + zero(TAT) * zero(TBT))) - # acc = -zero(Tacc) #Use this when KernelInterface - acc = @private Tacc 1 - @inbounds acc[1] = zero(Tacc) + Tacc = promote_type(R, typeof(zero(TAT) * zero(TBT) + zero(TAT) * zero(TBT))) + acc = zero(Tacc) # number of tiles depends on inner dimension - # NUM_TILES = div(K + TILE - 1, TILE) - @uniform NUM_TILES = cld(K, TILE) + NUM_TILES = cld(K, TILE) # loop over all tiles needed for this calculation for t in 0:(NUM_TILES - 1) - # Remove once moved to KernelInterface - gi = (grow - 1) * TILE + li - gj = (gcol - 1) * TILE + lj - # load inputs into tiles, with bounds checking for non-square matrices k0 = t * TILE ac = k0 + lj @@ -85,28 +76,20 @@ const TILE_DIM = 16 end # wait for all tiles to be loaded - @synchronize() - - # Remove once moved to KernelInterface - gi = (grow - 1) * TILE + li - gj = (gcol - 1) * TILE + lj + KI.barrier() @inbounds for k in 1:TILE - acc[1] = muladd(tile1[li, k], tile2[k, lj], acc[1]) + acc = muladd(tile1[li, k], tile2[k, lj], acc) end - @synchronize() + KI.barrier() end - # Remove once moved to KernelInterface - gi = (grow - 1) * TILE + li - gj = (gcol - 1) * TILE + lj - # save if inbounds if gi <= M && gj <= N - @inbounds C[gi, gj] = add(acc[1], C[gi, gj]) + @inbounds C[gi, gj] = add(acc, C[gi, gj]) end - # return + return end @inline function check_matmul_shapes(C, A, B, cA='N', cB='N') @@ -144,9 +127,7 @@ function generic_matmatmul!(C::AbstractArray{R}, cA::AbstractChar, cB::AbstractC workgroupsize=(TILE_DIM, TILE_DIM) numworkgroups=(cld(M, TILE_DIM), cld(N, TILE_DIM)) - # KI.@kernel KI.get_backend(C) workgroupsize=workgroupsize numworkgroups=numworkgroups coalesced_matmul_kernel!( - # C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM)) - coalesced_matmul_kernel!(get_backend(C), workgroupsize)( - C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM); ndrange=numworkgroups .* workgroupsize) + KI.@launch KI.get_backend(C) workgroupsize=workgroupsize numgroups=numworkgroups coalesced_matmul_kernel!( + C, A, B, add, Int(M), Int(N), Int(K), Val(cA), Val(cB), Val(TILE_DIM)) C end diff --git a/src/host/linalg.jl b/src/host/linalg.jl index 6e162b7f..631c0ba7 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -1178,7 +1178,7 @@ end ## reflect function LinearAlgebra.reflect!(x::AbstractGPUArray, y::AbstractGPUArray, c::Number, s::Number) - @kernel function reflect_kernel!(x, y, c, s) + @kernel function reflect_kernel!(x, y, c, s) i = @index(Global, Linear) @inbounds xi = x[i] @inbounds yi = y[i]