Support emitting a generated function

src/KernelAbstractions.jl

@@ -70,17 +70,26 @@ synchronize(dev)
 ```
 """
 macro kernel(expr)
-    return __kernel(expr, __source__, #=force_inbounds=# false, #=unsafe_indices=# false)
+    return __kernel(expr, __source__, __module__, #=force_inbounds=# false, #=unsafe_indices=# false, #=generated=# false)
 end
 
 """
     @kernel config function f(args) end
 
-This allows for two different configurations:
+This allows for the following configurations:
 
 1. `cpu={true, false}`: Disables code-generation of the CPU function. This relaxes semantics such that KernelAbstractions primitives can be used in non-kernel functions.
 2. `inbounds={false, true}`: Enables a forced `@inbounds` macro around the function definition in the case the user is using too many `@inbounds` already in their kernel. Note that this can lead to incorrect results, crashes, etc and is fundamentally unsafe. Be careful!
 3. `unsafe_indices={false, true}`: Disables the implicit validation of indices, users must avoid `@index(Global)`.
+4. `generated={false, true}`: Turns the kernel into a [generated function](https://docs.julialang.org/en/v1/manual/metaprogramming/#Generated-functions).
+   The kernel body is treated as a quoted expression, so `\$` interpolation is available and
+   `where`-parameters are bound to their values, e.g. to unroll a loop `\$N` times with `@unroll \$N for ...`.
+   This is meant for macros that need a literal, such as `@unroll \$N`, `Base.Cartesian.@nexprs \$N`
+   or `@ntuple \$N`; plain `where`-parameters are compile-time constants in every kernel already.
+   Configuration parameters must therefore be passed as types (`::Val{N}`) to be usable inside `\$`.
+   Inside `\$(...)` the argument names refer to the *types* of the arguments, not their values,
+   as in any generated function, and the body cannot contain closures, comprehensions or
+   generators (`x -> ...`, `do` blocks, `[f(i) for i in ...]`); use the Cartesian macros above instead.
 
 - [`@context`](@ref)
 
@@ -92,10 +101,11 @@ This allows for two different configurations:
 """
 macro kernel(ex...)
     if length(ex) == 1
-        return __kernel(ex[1], __source__, false, false)
+        return __kernel(ex[1], __source__, __module__, false, false, false)
     else
         unsafe_indices = false
         force_inbounds = false
+        generated = false
         for i in 1:(length(ex) - 1)
             if ex[i] isa Expr && ex[i].head == :(=) &&
                     ex[i].args[1] == :cpu && ex[i].args[2] isa Bool
@@ -106,17 +116,21 @@ macro kernel(ex...)
             elseif ex[i] isa Expr && ex[i].head == :(=) &&
                     ex[i].args[1] == :unsafe_indices && ex[i].args[2] isa Bool
                 unsafe_indices = ex[i].args[2]
+            elseif ex[i] isa Expr && ex[i].head == :(=) &&
+                    ex[i].args[1] == :generated && ex[i].args[2] isa Bool
+                generated = ex[i].args[2]
             else
                 error(
                     "Configuration should be of form:\n" *
                         "* `cpu=false`\n" *
                         "* `inbounds=true`\n" *
                         "* `unsafe_indices=true`\n" *
+                        "* `generated=true`\n" *
                         "got `", ex[i], "`",
                 )
             end
         end
-        return __kernel(ex[end], __source__, force_inbounds, unsafe_indices)
+        return __kernel(ex[end], __source__, __module__, force_inbounds, unsafe_indices, generated)
     end
 end
 

src/macros.jl

@@ -25,7 +25,7 @@ function unblock_lines(ex)
 end
 
 # XXX: Proper errors
-function __kernel(expr, __source__::LineNumberNode, force_inbounds = false, unsafe_indices = false)
+function __kernel(expr, __source__::LineNumberNode, __module__::Module, force_inbounds = false, unsafe_indices = false, generated = false)
     def = splitdef(expr)
     name = def[:name]
     args = def[:args]
@@ -46,6 +46,28 @@ function __kernel(expr, __source__::LineNumberNode, force_inbounds = false, unsa
     def_gpu = deepcopy(def)
     def_gpu[:name] = gpu_name = Symbol(:gpu_, name)
     transform_gpu!(def_gpu, constargs, force_inbounds, unsafe_indices)
+    if generated
+        # Turn the kernel into a generated function: the transformed body is
+        # quoted so that it is returned as an expression. Passing the quote
+        # through `macroexpand` (one level only, we do not want to expand the
+        # macros *inside* the quoted body here) lowers the `$` interpolations
+        # into the code that builds the expression at generation time.
+        body = macroexpand(__module__, Expr(:quote, def_gpu[:body]), recursive = false)
+        # Inference swallows any error thrown while generating (the kernel then
+        # merely infers to `Any`, which GPUCompiler reports as "kernel returns a
+        # value of type `Any`" without ever showing the cause), and there is no
+        # common launch path across backends where we could rethrow it. So
+        # catch the error here and hand it back through the return type
+        # instead, where GPUCompiler's validation prints it on every backend.
+        body = quote
+            try
+                $(check_generated)($(__module__), $body)
+            catch err
+                $(generated_error_body)(err)
+            end
+        end
+        def_gpu[:body] = Expr(:if, Expr(:generated), body, Expr(:meta, :generated_only))
+    end
     gpu_function = combinedef(def_gpu)
 
     # create constructor functions
@@ -66,6 +88,63 @@ function __kernel(expr, __source__::LineNumberNode, force_inbounds = false, unsa
     return Expr(:block, esc(gpu_function), esc(constructors))
 end
 
+"""
+    GeneratedKernelError{Msg}
+
+Marker type a `generated=true` kernel returns when its generator failed, carrying the
+error message in its type parameter. A kernel that returns this shows up in the
+`KernelError` GPUCompiler raises for kernels that return a value, which is the only
+channel through which a failure at generation time can be reported.
+"""
+struct GeneratedKernelError{Msg} end
+GeneratedKernelError(msg::AbstractString) = GeneratedKernelError{Symbol(msg)}()
+
+function Base.show(io::IO, ::Type{GeneratedKernelError{Msg}}) where {Msg}
+    return print(io, "KernelAbstractions.GeneratedKernelError(", repr(String(Msg)), ")")
+end
+
+# Runs inside the generator: makes sure the generated body is something Julia
+# accepts as the result of a generated function. Julia itself only rejects a
+# closure, comprehension or generator when lowering the body, which happens
+# outside the generator's `try` and thus can't be turned into a useful error
+# there, so check for them up front.
+function check_generated(mod::Module, body)
+    ex = macroexpand(mod, body)
+    MacroTools.postwalk(ex) do node
+        if isexpr(node, :->) || isexpr(node, :function) || isexpr(node, :do) ||
+                isexpr(node, :comprehension) || isexpr(node, :generator) ||
+                isexpr(node, :flatten) || (isexpr(node, :(=)) && isexpr(node.args[1], :call))
+            found = replace(string(MacroTools.striplines(node)), r"\s+" => " ")
+            error(
+                "the body of a `generated=true` kernel cannot contain a closure, " *
+                    "comprehension or generator (found `", found, "`). " *
+                    "Use `Base.Cartesian.@nexprs \$N` or `@ntuple \$N` instead.",
+            )
+        end
+        return node
+    end
+    return body
+end
+
+# Runs inside the generator, so it must not use code reflection: `showerror` for a
+# `MethodError` looks up candidate methods, which is forbidden there, so format
+# that one by hand and fall back to the bare exception type for anything else
+# that can't be shown.
+function generated_error_message(err)
+    if err isa MethodError
+        sig = join((a isa Type ? "::Type{$a}" : "::$(typeof(a))" for a in err.args), ", ")
+        return string("MethodError: no method matching ", err.f, "(", sig, ")")
+    end
+    msg = try
+        sprint(showerror, err)
+    catch
+        string(typeof(err))
+    end
+    return first(Base.split(msg, '\n'))
+end
+
+generated_error_body(err) = :(return $(GeneratedKernelError(generated_error_message(err))))
+
 # The easy case, transform the function for GPU execution
 # - mark constant arguments by applying `constify`.
 function transform_gpu!(def, constargs, force_inbounds, unsafe_indices)

test/unroll.jl

@@ -36,6 +36,39 @@ end
     end
 end
 
+# `generated=true` makes the kernel a generated function, so that the `where`
+# parameter `N` can be interpolated into `@unroll $N`, which requires a literal.
+@kernel generated = true function kernel_unroll_generated!(a, ::Val{N}) where {N}
+    @unroll $N for i in 1:5
+        @inbounds a[i] = i * $N
+    end
+end
+
+# `generated=true` composes with the other body transformations: `@Const`,
+# `@localmem`, `@synchronize` and `inbounds=true` all round-trip through the quote.
+@kernel generated = true inbounds = true function kernel_generated_transforms!(a, @Const(b), ::Val{N}) where {N}
+    tile = @localmem Float32 (N,)
+    I = @index(Global, Linear)
+    i = @index(Local, Linear)
+    @unroll $N for k in 1:N
+        tile[k] = b[k] * $N
+    end
+    @synchronize
+    a[I] = tile[i]
+end
+
+# Errors while generating are reported through the return type rather than being
+# swallowed as `Any`; the message must carry the original error.
+@kernel generated = true function kernel_generated_closure!(a)
+    I = @index(Global)
+    @inbounds a[I] = sum(x -> x, 1:$(2))
+end
+
+@kernel generated = true function kernel_generated_badinterp!(a)
+    I = @index(Global)
+    @inbounds a[I] = $(length(a))
+end
+
 function unroll_testsuite(backend, ArrayT)
     a = ArrayT(zeros(Float32, 5))
     kernel! = kernel_unroll!(backend(), 1, 1)
@@ -44,5 +77,39 @@ function unroll_testsuite(backend, ArrayT)
     kernel2! = kernel_unroll2!(backend(), 1, 1)
     kernel2!(a)
     synchronize(backend())
+
+    a = ArrayT(zeros(Float32, 5))
+    kernel3! = kernel_unroll_generated!(backend(), 1, 1)
+    kernel3!(a, Val(2))
+    synchronize(backend())
+    @test Array(a) == Float32[2, 4, 6, 8, 10]
+
+    a = ArrayT(zeros(Float32, 4))
+    b = ArrayT(Float32[1, 2, 3, 4])
+    kernel4! = kernel_generated_transforms!(backend(), 4, 4)
+    kernel4!(a, b, Val(4))
+    synchronize(backend())
+    @test Array(a) == Float32[4, 8, 12, 16]
+
+    a = ArrayT(zeros(Float32, 2))
+    err = try
+        kernel_generated_closure!(backend(), 2)(a; ndrange = 2)
+        synchronize(backend())
+        nothing
+    catch e
+        sprint(showerror, e)
+    end
+    @test occursin("GeneratedKernelError", err)
+    @test occursin("cannot contain a closure", err)
+
+    err = try
+        kernel_generated_badinterp!(backend(), 2)(a; ndrange = 2)
+        synchronize(backend())
+        nothing
+    catch e
+        sprint(showerror, e)
+    end
+    @test occursin("GeneratedKernelError", err)
+    @test occursin("MethodError: no method matching length(::Type{", err)
     return
 end