diff --git a/Project.toml b/Project.toml index 1bd6c033..88649a38 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" Adapt = "4.7.3" GPUArraysCore = "= 0.2.1" 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/src/host/gemm.jl b/src/host/gemm.jl index d018806d..cd37f757 100644 --- a/src/host/gemm.jl +++ b/src/host/gemm.jl @@ -1,5 +1,5 @@ using LinearAlgebra: MulAddMul, wrapper_char, _unwrap - +import KernelInterface as KI ## Matmul kernel # A[i, j] of a symmetric/hermitian matrix with only one triangle stored @@ -32,43 +32,37 @@ end end end +# additive identity to start accumulating from: -0.0, not 0.0, for IEEE floats, so that a +# sum of -0.0 products stays -0.0. other element types needn't support unary minus. +@inline accumulator_zero(::Type{T}) where {T <: Number} = -zero(T) +@inline accumulator_zero(::Type{T}) where {T} = zero(T) + 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) - - # li, lj, _ = KI.get_local_id() - li, lj = @index(Local, NTuple) + TAT = eltype(A); TBT = eltype(B); R = eltype(C) - # 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) + # Revisit potentially padding if value of `TILE_DIM` ever changed + 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 = accumulator_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 +79,21 @@ 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]) + # skip the zero padding past K: adding its +0.0 products would turn a -0.0 into +0.0 + @inbounds for k in 1:min(TILE, K - k0) + 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 +131,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 3c6788e2..c80c4469 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -1182,7 +1182,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]