diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 276962365..6f88988d3 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -658,6 +658,21 @@ function __workitems_iterspace end end end +# Whether all work-items of the current workgroup are part of the ndrange. `@kernel` runs the +# body of such a workgroup without testing `__validindex` for every work-item. Backends for +# which the masked body is slow (e.g., because their compiler vectorizes over the work-items of +# a workgroup, as PoCL does) can override this to return `__fullgroup_check(ctx)`; the default +# never takes that path, so the unmasked copy of the body is removed when compiling. +@inline __fullgroup(ctx) = false + +@inline function __fullgroup_check(ctx) + if __dynamic_checkbounds(ctx) + return full_group(ctx, index_launch(ctx)) + else + return true + end +end + include("macros.jl") include("spawn.jl") include("foreach_index.jl") diff --git a/src/launch.jl b/src/launch.jl index 37667a75b..7ef14bf07 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -274,6 +274,18 @@ else global_cartesian(ctx, launch, iterspace, ndrange) in ndrange end +# Whether the current workgroup lies entirely inside the ndrange. This is the same for all +# work-items of the workgroup. +@inline full_group(ctx, launch::Launch, iterspace, ndrange) = +if builtin(iterspace) && ndrange isa CartesianIndices + T = index_type(launch) + groupsize = narrow(T, size(workitems(iterspace))) + reduce(&, map((g, w, n) -> g * w <= n, group_position(ctx, launch), groupsize, narrow(T, size(ndrange))); init = true) +else + false +end +@inline full_group(ctx, launch::Launch) = full_group(ctx, launch, __iterspace(ctx), __ndrange(ctx)) + # `@index` entry points, see `__index_*` @inline local_linear(ctx, ::LinearLaunch{T}) where {T} = widen_index(KI.get_local_id(T).x) @inline group_linear(ctx, ::LinearLaunch{T}) where {T} = widen_index(KI.get_group_id(T).x) diff --git a/src/macros.jl b/src/macros.jl index 3d19333c8..cd12122bd 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -178,6 +178,7 @@ function transform_gpu!(def, constargs, force_inbounds, unsafe_indices) new_stmts = Any[] body = MacroTools.flatten(def[:body]) if !unsafe_indices + push!(new_stmts, :(__full_group__ = $__fullgroup(__ctx__))) push!(new_stmts, :(__active_lane__ = $__validindex(__ctx__))) end if force_inbounds @@ -358,7 +359,14 @@ function emit(loop) stmts = Any[] append!(stmts, loop.allocations) - push!(stmts, Expr(:if, :__active_lane__, Expr(:block, loop.stmts...))) + # The body is emitted twice: without a mask for workgroups that lie entirely inside the + # ndrange (the branch is the same for all their work-items), and masked for the others. + push!( + stmts, Expr( + :if, :__full_group__, Expr(:block, loop.stmts...), + Expr(:elseif, :__active_lane__, Expr(:block, deepcopy(loop.stmts)...)) + ) + ) if loop.terminated_in_sync loop.sync_line === nothing || push!(stmts, loop.sync_line) push!(stmts, :($__synchronize())) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 4d5df6726..77393d345 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -237,6 +237,13 @@ end end +## Bounds checks + +# PoCL vectorizes over the work-items of a workgroup, and masks the whole body of a kernel +# whose work-items may be outside the ndrange. Run full workgroups without the mask. +@device_override @inline KA.__fullgroup(ctx) = KA.__fullgroup_check(ctx) + + ## Synchronization and Printing @device_override @inline function KI.barrier() diff --git a/test/codegen_checks.jl b/test/codegen_checks.jl index 6db27c38c..e9dd0b2a2 100644 --- a/test/codegen_checks.jl +++ b/test/codegen_checks.jl @@ -122,6 +122,30 @@ end end end + # On PoCL, the body of a kernel with a dynamic ndrange is emitted twice: unmasked for the + # workgroups that lie inside the ndrange, and masked for a partial one. + @testset "full workgroups" begin + @test @filecheck begin + @check "define spir_kernel void @{{.*}}gpu_codegen_mul2_inbounds" + @check "store float" + @check "store float" + @check_not "store float" + @check "ret void" + @device_code_llvm debuginfo = :none codegen_mul2_inbounds(backend, 16)(A, ndrange = 64) + KernelAbstractions.synchronize(backend) + end + + # with the workgroups known to lie inside the ndrange, only the unmasked body is left + @test @filecheck begin + @check "define spir_kernel void @{{.*}}gpu_codegen_mul2_inbounds" + @check "store float" + @check_not "store float" + @check "ret void" + @device_code_llvm debuginfo = :none codegen_mul2_inbounds(backend, 16, 64)(A) + KernelAbstractions.synchronize(backend) + end + end + # `@localmem` becomes a module-level allocation in the SPIR-V workgroup address space # (3), which the kernel reads and writes directly. The two accesses are `@check_dag` # because LLVM is free to emit the basic blocks in any order. @@ -181,13 +205,14 @@ end end end - # `@print` lowers to a single variadic printf call, not to one call per argument. + # `@print` lowers to a single variadic printf call, not to one call per argument. (With a + # static ndrange, so that the body isn't emitted twice, see "full workgroups" below.) @testset "print" begin @test @filecheck begin @check "define spir_kernel void @{{.*}}gpu_codegen_print" @check "@printf" @check_not "@printf" - @device_code_llvm debuginfo = :none codegen_print(backend, 16)(ndrange = 16) + @device_code_llvm debuginfo = :none codegen_print(backend, 16, 16)() KernelAbstractions.synchronize(backend) end end