Skip to content

Commit 39cf28b

Browse files
committed
avoid boxing when @kernel is used as a closure
This allows uses of `@kernel` inside of functions without running into JuliaLang/julia#53295. One thing I have been wondering though: Is the `@isdefined` check here really necessary? In global scope it should just override the previous definition and for closures these kind of if statements are ignored anyways.
1 parent b8b53da commit 39cf28b

File tree

2 files changed

+21
-4
lines changed

2 files changed

+21
-4
lines changed

src/macros.jl

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,10 @@ function __kernel(expr, generate_cpu = true, force_inbounds = false, unsafe_indi
5050
gpu_function = combinedef(def_gpu)
5151

5252
# create constructor functions
53+
_name = Symbol(:_, name)
5354
constructors = quote
5455
if $(name isa Symbol ? :(!@isdefined($name)) : true)
55-
Core.@__doc__ $name(dev) = $name(dev, $DynamicSize(), $DynamicSize())
56-
$name(dev, size) = $name(dev, $StaticSize(size), $DynamicSize())
57-
$name(dev, size, range) = $name(dev, $StaticSize(size), $StaticSize(range))
58-
function $name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size}
56+
function $_name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size}
5957
if $isgpu(dev)
6058
return $construct(dev, sz, range, $gpu_name)
6159
else
@@ -66,6 +64,10 @@ function __kernel(expr, generate_cpu = true, force_inbounds = false, unsafe_indi
6664
end
6765
end
6866
end
67+
Core.@__doc__ $name(dev) = $_name(dev, $DynamicSize(), $DynamicSize())
68+
$name(dev, size) = $_name(dev, $StaticSize(size), $DynamicSize())
69+
$name(dev, size, range) = $_name(dev, $StaticSize(size), $StaticSize(range))
70+
$name(dev, size::$_Size, range::$_Size) = $_name(dev, size, range)
6971
end
7072
end
7173

test/test.jl

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -349,5 +349,20 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk
349349
end
350350
end
351351

352+
@testset "`@kernel` as a closure" begin
353+
function foo()
354+
@kernel function kernel(A)
355+
i = @index(Global)
356+
A[i] = 1
357+
end
358+
return kernel
359+
end
360+
ftypes = fieldtypes(typeof(foo()))
361+
@test !any(T -> T <: Core.Box, ftypes)
362+
@test all(ftypes) do T
363+
!any(T -> T <: Core.Box, fieldtypes(T))
364+
end
365+
end
366+
352367
return
353368
end

0 commit comments

Comments
 (0)