Skip to content

Commit 434d55a

Browse files
authored
avoid boxing when @kernel is used as a closure (#625)
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 434d55a

File tree

3 files changed

+22
-5
lines changed

3 files changed

+22
-5
lines changed

Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
name = "KernelAbstractions"
22
uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
33
authors = ["Valentin Churavy <[email protected]> and contributors"]
4-
version = "0.9.36"
4+
version = "0.9.37"
55

66
[deps]
77
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

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)