From 39cf28b9be35fb30d6d2f2d2a97e61b0a234cadd Mon Sep 17 00:00:00 2001
From: Simeon David Schaub <simeon@schaub.rocks>
Date: Fri, 4 Jul 2025 10:53:12 +0200
Subject: [PATCH] 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.
---
 src/macros.jl | 10 ++++++----
 test/test.jl  | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/macros.jl b/src/macros.jl
index c83d60308..86d5b0bcb 100644
--- a/src/macros.jl
+++ b/src/macros.jl
@@ -50,12 +50,10 @@ function __kernel(expr, generate_cpu = true, force_inbounds = false, unsafe_indi
     gpu_function = combinedef(def_gpu)
 
     # create constructor functions
+    _name = Symbol(:_, name)
     constructors = quote
         if $(name isa Symbol ? :(!@isdefined($name)) : true)
-            Core.@__doc__ $name(dev) = $name(dev, $DynamicSize(), $DynamicSize())
-            $name(dev, size) = $name(dev, $StaticSize(size), $DynamicSize())
-            $name(dev, size, range) = $name(dev, $StaticSize(size), $StaticSize(range))
-            function $name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size}
+            function $_name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size}
                 if $isgpu(dev)
                     return $construct(dev, sz, range, $gpu_name)
                 else
@@ -66,6 +64,10 @@ function __kernel(expr, generate_cpu = true, force_inbounds = false, unsafe_indi
                     end
                 end
             end
+            Core.@__doc__ $name(dev) = $_name(dev, $DynamicSize(), $DynamicSize())
+            $name(dev, size) = $_name(dev, $StaticSize(size), $DynamicSize())
+            $name(dev, size, range) = $_name(dev, $StaticSize(size), $StaticSize(range))
+            $name(dev, size::$_Size, range::$_Size) = $_name(dev, size, range)
         end
     end
 
diff --git a/test/test.jl b/test/test.jl
index 4a4f667ab..7de6da0bb 100644
--- a/test/test.jl
+++ b/test/test.jl
@@ -349,5 +349,20 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk
         end
     end
 
+    @testset "`@kernel` as a closure" begin
+        function foo()
+            @kernel function kernel(A)
+                i = @index(Global)
+                A[i] = 1
+            end
+            return kernel
+        end
+        ftypes = fieldtypes(typeof(foo()))
+        @test !any(T -> T <: Core.Box, ftypes)
+        @test all(ftypes) do T
+            !any(T -> T <: Core.Box, fieldtypes(T))
+        end
+    end
+
     return
 end