Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Tag dispatch #78

Closed
wants to merge 2 commits into from
Closed

Tag dispatch #78

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
abe26c3
Improve block and thread calculations and invoke only if in range
PhilipFackler Apr 18, 2024
025175b
Add tag dispatch to separate implementations
PhilipFackler Apr 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions ext/JACCAMDGPU/JACCAMDGPU.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,17 @@ using JACC, AMDGPU

include("array.jl")

function JACC.parallel_for(N::I, f::F, x...) where {I <: Integer, F <: Function}
struct AMDGPUTag end

function JACC.parallel_for(::AMDGPUTag, N::I, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 512
threads = min(N, numThreads)
blocks = ceil(Int, N / threads)
@roc groupsize = threads gridsize = blocks _parallel_for_amdgpu(f, x...)
AMDGPU.synchronize()
end

function JACC.parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_for(::AMDGPUTag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 16
Mthreads = min(M, numThreads)
Nthreads = min(N, numThreads)
Expand All @@ -22,7 +24,7 @@ function JACC.parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer,
AMDGPU.synchronize()
end

function JACC.parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_reduce(::AMDGPUTag, N::I, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 512
threads = min(N, numThreads)
blocks = ceil(Int, N / threads)
Expand All @@ -36,7 +38,7 @@ function JACC.parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Functi

end

function JACC.parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_reduce(::AMDGPUTag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 16
Mthreads = min(M, numThreads)
Nthreads = min(N, numThreads)
Expand Down Expand Up @@ -303,6 +305,7 @@ end

function __init__()
const JACC.Array = AMDGPU.ROCArray{T, N} where {T, N}
const JACC.Tag = AMDGPUTag
end

end # module JACCAMDGPU
34 changes: 24 additions & 10 deletions ext/JACCCUDA/JACCCUDA.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,25 @@ using JACC, CUDA
# overloaded array functions
include("array.jl")

function JACC.parallel_for(N::I, f::F, x...) where {I <: Integer, F <: Function}
maxPossibleThreads = attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X)
threads = min(N, maxPossibleThreads)
blocks = ceil(Int, N / threads)
CUDA.@sync @cuda threads = threads blocks = blocks _parallel_for_cuda(f, x...)
struct CUDATag end

function JACC.parallel_for(::CUDATag, N::I, f::F, x...) where {I <: Integer, F <: Function}
parallel_args = (N, f, x...)
parallel_kargs = cudaconvert.(parallel_args)
parallel_tt = Tuple{Core.Typeof.(parallel_kargs)...}
parallel_kernel = cufunction(_parallel_for_cuda, parallel_tt)
maxPossibleThreads = CUDA.maxthreads(parallel_kernel)
threads = min(N, maxPossibleThreads)
blocks = ceil(Int, N / threads)
parallel_kernel(parallel_kargs...; threads=threads, blocks=blocks)

# maxPossibleThreads = attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X)
# threads = min(N, maxPossibleThreads)
# blocks = ceil(Int, N / threads)
# CUDA.@sync @cuda threads = threads blocks = blocks _parallel_for_cuda(f, x...)
end

function JACC.parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_for(::CUDATag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 16
Mthreads = min(M, numThreads)
Nthreads = min(N, numThreads)
Expand All @@ -21,7 +32,7 @@ function JACC.parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer,
CUDA.@sync @cuda threads = (Mthreads, Nthreads) blocks = (Mblocks, Nblocks) _parallel_for_cuda_MN(f, x...)
end

function JACC.parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_reduce(::CUDATag, N::I, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 512
threads = min(N, numThreads)
blocks = ceil(Int, N / threads)
Expand All @@ -33,7 +44,7 @@ function JACC.parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Functi
end


function JACC.parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_reduce(::CUDATag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
numThreads = 16
Mthreads = min(M, numThreads)
Nthreads = min(N, numThreads)
Expand All @@ -46,9 +57,11 @@ function JACC.parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integ
return rret
end

function _parallel_for_cuda(f, x...)
function _parallel_for_cuda(N, f, x...)
i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
f(i, x...)
if i <= N
f(i, x...)
end
return nothing
end

Expand Down Expand Up @@ -298,6 +311,7 @@ end

function __init__()
const JACC.Array = CUDA.CuArray{T, N} where {T, N}
const JACC.Tag = CUDATag
end

end # module JACCCUDA
11 changes: 7 additions & 4 deletions ext/JACCONEAPI/JACCONEAPI.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,17 @@ module JACCONEAPI

using JACC, oneAPI

function JACC.parallel_for(N::I, f::F, x...) where {I <: Integer, F <: Function}
struct oneAPITag end

function JACC.parallel_for(::oneAPITag, N::I, f::F, x...) where {I <: Integer, F <: Function}
#maxPossibleItems = oneAPI.oneL0.compute_properties(device().maxTotalGroupSize)
maxPossibleItems = 256
items = min(N, maxPossibleItems)
groups = ceil(Int, N / items)
oneAPI.@sync @oneapi items = items groups = groups _parallel_for_oneapi(f, x...)
end

function JACC.parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_for(::oneAPITag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
maxPossibleItems = 16
Mitems = min(M, maxPossibleItems)
Nitems = min(N, maxPossibleItems)
Expand All @@ -20,7 +22,7 @@ function JACC.parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer,
oneAPI.@sync @oneapi items = (Mitems, Nitems) groups = (Mgroups, Ngroups) _parallel_for_oneapi_MN(f, x...)
end

function JACC.parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_reduce(::oneAPITag, N::I, f::F, x...) where {I <: Integer, F <: Function}
numItems = 256
items = min(N, numItems)
groups = ceil(Int, N / items)
Expand All @@ -31,7 +33,7 @@ function JACC.parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Functi
return rret
end

function JACC.parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
function JACC.parallel_reduce(::oneAPITag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
numItems = 16
Mitems = min(M, numItems)
Nitems = min(N, numItems)
Expand Down Expand Up @@ -295,6 +297,7 @@ end

function __init__()
const JACC.Array = oneAPI.oneArray{T, N} where {T, N}
const JACC.Tag = oneAPITag
end

end # module JACCONEAPI
34 changes: 25 additions & 9 deletions src/JACC.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,33 +1,41 @@
__precompile__(false)
# __precompile__(false)
module JACC

import Atomix: @atomic
# module to set back end preferences
include("JACCPreferences.jl")
include("helper.jl")
# overloaded array functions
include("array.jl")

export Array, @atomic
export Array
export parallel_for

global Array
global Tag

function parallel_for(N::I, f::F, x...) where {I <: Integer, F <: Function}
struct ThreadsTag end

function parallel_for(::ThreadsTag, N::I, f::F, x...) where {I <: Integer, F <: Function}
@maybe_threaded for i in 1:N
f(i, x...)
end
end

function parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
@inline function parallel_for(N::I, f::F, x...) where {I <: Integer, F <: Function}
parallel_for(Tag(), N, f, x...)
end

function parallel_for(::ThreadsTag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
@maybe_threaded for j in 1:N
for i in 1:M
f(i, j, x...)
end
end
end

function parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Function}
@inline function parallel_for((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
parallel_for(Tag(), (M, N), f, x...)
end

function parallel_reduce(::ThreadsTag, N::I, f::F, x...) where {I <: Integer, F <: Function}
tmp = zeros(Threads.nthreads())
ret = zeros(1)
@maybe_threaded for i in 1:N
Expand All @@ -39,7 +47,11 @@ function parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Function}
return ret
end

function parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
@inline function parallel_reduce(N::I, f::F, x...) where {I <: Integer, F <: Function}
parallel_reduce(Tag(), N, f, x...)
end

function parallel_reduce(::ThreadsTag, (M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
tmp = zeros(Threads.nthreads())
ret = zeros(1)
@maybe_threaded for j in 1:N
Expand All @@ -53,9 +65,13 @@ function parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F
return ret
end

@inline function parallel_reduce((M, N)::Tuple{I, I}, f::F, x...) where {I <: Integer, F <: Function}
parallel_reduce(Tag(), (M, N), f, x...)
end

function __init__()
const JACC.Array = Base.Array{T, N} where {T, N}
const JACC.Tag = ThreadsTag
end


Expand Down
Loading