Skip to content

Fix launch parameter heuristic exceeding the maximum number of threads per block #216

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

Merged
omlins merged 1 commit into from
Aug 20, 2026
+27 −9
Merged

Fix launch parameter heuristic exceeding the maximum number of threads per block #216

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
16 changes: 8 additions & 8 deletions src/ParallelKernel/parallel.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -485,7 +485,7 @@ function handle_inverses(body::Expr)
end

function handle_padding(caller::Module, body::Expr, padding::Bool, indices; handle_view_accesses::Bool=true, handle_indexing::Bool=true, dir_handling::Bool=true, delay_dir_handling::Bool=false)
if (handle_indexing)
if (handle_indexing)
body = substitute_indices_inn(body, padding)
if (dir_handling) body = substitute_indices_dir(caller, body, padding; delay_handling=delay_dir_handling) end
body = substitute_firstlastindex(caller, body, padding)
Expand Down Expand Up @@ -519,7 +519,7 @@ function substitute_indices_dir(caller::Module, expr::Expr, padding::Bool; delay
if @capture(ex, A_[indices_expr__]) && any(map(inexpr_walk, indices_expr, INDICES_DIR))
A_parent = promote_to_parent(A)
ex = substitute(ex, NamedTuple{INDICES_DIR}(
((A_parent==B_parent) ? ix : :($ix - (size($B_parent, 1) > size($A_parent, 1))),
((A_parent==B_parent) ? ix : :($ix - (size($B_parent, 1) > size($A_parent, 1))),
(A_parent==B_parent) ? iy : :($iy - (size($B_parent, 2) > size($A_parent, 2))),
(A_parent==B_parent) ? iz : :($iz - (size($B_parent, 3) > size($A_parent, 3))))
); inQuoteNode=true)
Expand Down Expand Up @@ -558,7 +558,7 @@ end

function substitute_firstlastindex(caller::Module, body::Expr, padding::Bool)
return postwalk(body) do ex
if @capture(ex, f_(args__))
if @capture(ex, f_(args__))
if (f == :firstindex) return _firstindex(caller, args..., padding)
elseif (f == :lastindex) return _lastindex(caller, args..., padding)
else return ex
Expand Down Expand Up @@ -804,9 +804,9 @@ end

function compute_nthreads(maxsize; nthreads_x_max=NTHREADS_X_MAX, nthreads_max=NTHREADS_MAX, flatdim=0) # This is a heuristic, which results in (32,8,1) threads, except if maxsize[1] < 32 or maxsize[2] < 8.
maxsize = promote_maxsize(maxsize)
nthreads_x = min(nthreads_x_max, (flatdim==1) ? 1 : maxsize[1])
nthreads_y = min(ceil(Int,nthreads_max/nthreads_x), (flatdim==2) ? 1 : maxsize[2])
nthreads_z = min(ceil(Int,nthreads_max/(nthreads_x*nthreads_y)), (flatdim==3) ? 1 : maxsize[3])
nthreads_x = min(nthreads_x_max, (flatdim==1) ? 1 : maxsize[1])
nthreads_y = min(max(floor(Int,nthreads_max/nthreads_x), 1), (flatdim==2) ? 1 : maxsize[2])
nthreads_z = min(max(floor(Int,nthreads_max/(nthreads_x*nthreads_y)), 1), (flatdim==3) ? 1 : maxsize[3]) # NOTE: the thread budget of each dimension is rounded down, because rounding it up can make the total number of threads exceed nthreads_max whenever a previous dimension was clamped to a smaller maxsize.
return (nthreads_x, nthreads_y , nthreads_z)
end

Expand Down Expand Up @@ -892,7 +892,7 @@ function create_gpu_or_xpu_call(package::Symbol, nblocks::Union{Symbol,Expr}, nt
else @ModuleInternalError("unsupported GPU package (obtained: $package).")
end
if !isnothing(shmem_expr)
backend_kwargs_expr = (backend_kwargs_expr..., shmem_expr)
backend_kwargs_expr = (backend_kwargs_expr..., shmem_expr)
end
end
if (package == PKG_CUDA) return :( CUDA.@cuda blocks=$nblocks threads=$nthreads stream=$stream $(backend_kwargs_expr...) $kernelcall; $synccall )
Expand Down Expand Up @@ -933,4 +933,4 @@ function default_stream(package)
elseif (package == PKG_KERNELABSTRACTIONS) return :(nothing)
else @ModuleInternalError("unsupported GPU package (obtained: $package).")
end
end
end
20 changes: 19 additions & 1 deletion test/ParallelKernel/test_parallel.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ using ParallelStencil.ParallelKernel
import ParallelStencil.ParallelKernel.AD
import ParallelStencil.ParallelKernel: @reset_parallel_kernel, @is_initialized, SUPPORTED_PACKAGES, PKG_CUDA, PKG_AMDGPU, PKG_METAL, PKG_THREADS, PKG_POLYESTER, PKG_KERNELABSTRACTIONS, INDICES, ARRAYTYPES, FIELDTYPES, SCALARTYPES
import ParallelStencil.ParallelKernel: @require, @prettystring, @gorgeousstring, @isgpu, @iscpu, interpolate, @select_hardware, @current_hardware, handle
import ParallelStencil.ParallelKernel: checkargs_parallel, checkargs_parallel_indices, parallel_indices, maxsize
import ParallelStencil.ParallelKernel: checkargs_parallel, checkargs_parallel_indices, parallel_indices, maxsize, compute_nthreads, NTHREADS_X_MAX_AMDGPU, NTHREADS_MAX
using ParallelStencil.ParallelKernel.Exceptions
TEST_PACKAGES = SUPPORTED_PACKAGES
@static if PKG_CUDA in TEST_PACKAGES
Expand Down Expand Up @@ -219,6 +219,24 @@ eval(:(
@test maxsize((x=8, y=[9 9; 9 9; 9 9]), [7 7 7; 7 7 7]) == (3, 3, 1)
@test maxsize(BitstypeStruct(5, 6.0), 8, (x=[9 9; 9 9; 9 9], y=[9 9; 9 9; 9 9]), (x=[7 7 7; 7 7 7], y=[7 7 7; 7 7 7])) == (3, 3, 1)
end;
@testset "compute_nthreads" begin
@test compute_nthreads(256) == (32, 1, 1)
@test compute_nthreads((256, 256)) == (32, 8, 1)
@test compute_nthreads((256, 256, 256)) == (32, 8, 1)
@test compute_nthreads((256, 256, 256); nthreads_x_max=NTHREADS_X_MAX_AMDGPU) == (64, 4, 1)
@test compute_nthreads((8, 8, 8)) == (8, 8, 4)
@test compute_nthreads((256, 256, 256); nthreads_max=128, flatdim=3) == (32, 4, 1)
@testset "nthreads never exceeds nthreads_max" begin # NOTE: rounding the thread budget of each dimension up used to make the total number of threads exceed nthreads_max whenever a previous dimension was clamped to a smaller maxsize, resulting in a launch failure (e.g. maxsize=(256,5,128) resulted in (32,5,2)=320 threads).
@test compute_nthreads((256, 5, 128)) == (32, 5, 1)
@test compute_nthreads((256, 3, 128)) == (32, 3, 2)
@test compute_nthreads((10, 100, 100)) == (10, 25, 1)
maxsizes = ((i, j, k) for i in 1:70, j in 1:70, k in 1:70) # NOTE: `ms` is used below instead of `maxsize` in order not to shadow the imported function `maxsize`.
@test all(prod(compute_nthreads(ms)) <= NTHREADS_MAX for ms in maxsizes)
@test all(prod(compute_nthreads(ms; nthreads_x_max=NTHREADS_X_MAX_AMDGPU)) <= NTHREADS_MAX for ms in maxsizes)
@test all(prod(compute_nthreads(ms; nthreads_max=128, flatdim=loopdim)) <= 128 for ms in maxsizes, loopdim in 1:3)
@test all(all(compute_nthreads(ms) .>= 1) for ms in maxsizes)
end;
end;
end;
@static if $package != $PKG_POLYESTER # Enzyme does not support Polyester.
@testset "@parallel ∇" begin
Expand Down
Loading