From e124181fde1bbcad96c6c27d82e7e20ff57d67f0 Mon Sep 17 00:00:00 2001 From: Anton Smirnov Date: Fri, 18 Oct 2024 12:07:28 +0300 Subject: [PATCH 1/3] Allow return statements for GPU-only kernels --- src/macros.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/macros.jl b/src/macros.jl index 55b4bab40..a511758dc 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -14,7 +14,9 @@ function __kernel(expr, generate_cpu = true, force_inbounds = false) def = splitdef(expr) name = def[:name] args = def[:args] - find_return(expr) && error("Return statement not permitted in a kernel function $name") + generate_cpu && find_return(expr) && error( + "Return statement not permitted in a kernel function $name", + ) constargs = Array{Bool}(undef, length(args)) for (i, arg) in enumerate(args) From 860c71f87d0b1a1ff4b2862d07932bcc10f160f3 Mon Sep 17 00:00:00 2001 From: Anton Smirnov Date: Wed, 23 Oct 2024 23:54:50 +0300 Subject: [PATCH 2/3] Add tests --- test/test.jl | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test/test.jl b/test/test.jl index 8977231b7..7bcaa6f3f 100644 --- a/test/test.jl +++ b/test/test.jl @@ -306,4 +306,22 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @test size(KernelAbstractions.zeros(backend, Float32, 0, 9)) == (0, 9) end + @kernel cpu=false function gpu_return_kernel!(x) + i = @index(Global) + if i ≤ (length(x) ÷ 2) + x[i] = 1 + return + end + end + @testset "GPU kernel return statement" begin + if !(Backend() isa CPU) + A = KernelAbstractions.zeros(Backend(), Int64, 1024) + gpu_return_kernel!(Backend())(A; ndrange=length(A)) + synchronize(Backend()) + Ah = Array(A) + @test all(a -> a == 1, @view(Ah[1:length(A) ÷ 2])) + @test all(a -> a == 0, @view(Ah[length(A) ÷ 2 + 1:end])) + end + end + end From b3f66760325ea4f16c66ab465fd09ed14d525a92 Mon Sep 17 00:00:00 2001 From: Anton Smirnov Date: Thu, 24 Oct 2024 00:00:46 +0300 Subject: [PATCH 3/3] Refactor --- test/test.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/test.jl b/test/test.jl index 7bcaa6f3f..337af10d4 100644 --- a/test/test.jl +++ b/test/test.jl @@ -306,7 +306,7 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @test size(KernelAbstractions.zeros(backend, Float32, 0, 9)) == (0, 9) end - @kernel cpu=false function gpu_return_kernel!(x) + @kernel cpu = false function gpu_return_kernel!(x) i = @index(Global) if i ≤ (length(x) ÷ 2) x[i] = 1 @@ -316,11 +316,11 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @testset "GPU kernel return statement" begin if !(Backend() isa CPU) A = KernelAbstractions.zeros(Backend(), Int64, 1024) - gpu_return_kernel!(Backend())(A; ndrange=length(A)) + gpu_return_kernel!(Backend())(A; ndrange = length(A)) synchronize(Backend()) Ah = Array(A) - @test all(a -> a == 1, @view(Ah[1:length(A) ÷ 2])) - @test all(a -> a == 0, @view(Ah[length(A) ÷ 2 + 1:end])) + @test all(a -> a == 1, @view(Ah[1:(length(A) ÷ 2)])) + @test all(a -> a == 0, @view(Ah[(length(A) ÷ 2 + 1):end])) end end