diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index d7ca0ac8a..7c2c737a1 100644 --- a/src/KernelAbstractions.jl +++ b/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 diff --git a/src/macros.jl b/src/macros.jl index d0b05e707..5ba32bb2d 100644 --- a/src/macros.jl +++ b/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) diff --git a/test/unroll.jl b/test/unroll.jl index c9a64466d..1cd477a02 100644 --- a/test/unroll.jl +++ b/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