Run full workgroups without the bounds mask on the CPU backend

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")

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)

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()))

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()

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