From 208fb93504ad805ccd6a606acf9f343b4a42e5d5 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 5 Nov 2026 18:07:37 +0100 Subject: [PATCH 1/2] KernelInterface: inline the generated `argument_types` `@inline` on a generated function only marks the body for inlining when the generator falls through to it, not when it returns it explicitly, as Runic formats it. Move the `inline_body` helper that the shuffles use to the utilities, and use it for `argument_types` too. --- lib/KernelInterface/src/device.jl | 4 ---- lib/KernelInterface/src/launch.jl | 2 +- lib/KernelInterface/src/utils.jl | 6 ++++++ 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 76792eafc..15528ddf9 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -411,10 +411,6 @@ shuffle_as_words(T) = T !== UInt32 && sizeof(T) in (1, 2, 4, 8, 16) const word_types = (UInt8, UInt16, UInt32, UInt64, UInt128) word_type(T) = word_types[trailing_zeros(sizeof(T)) + 1] -# The generated functions mark their bodies for inlining: `@inline` on a generated function -# doesn't apply to a body that the generator returns with an explicit `return`. -inline_body(ex) = Expr(:block, Expr(:meta, :inline), ex) - # Shuffle a primitive value as unsigned words: values of up to 4 bytes are zero-extended to a # `UInt32`. Other 8-byte values are shuffled as a `UInt64`, which the backend may support # natively, and a `UInt64` it doesn't is split into two `UInt32` words. 16-byte values are diff --git a/lib/KernelInterface/src/launch.jl b/lib/KernelInterface/src/launch.jl index 86e4babd2..29e3c9735 100644 --- a/lib/KernelInterface/src/launch.jl +++ b/lib/KernelInterface/src/launch.jl @@ -437,7 +437,7 @@ function kernel_function end # isn't type stable for 32 or more elements @inline @generated function argument_types(backend, args::Tuple) types = (:(Core.Typeof(argconvert(backend, args[$i]))) for i in 1:fieldcount(args)) - return :(Tuple{$(types...)}) + return inline_body(:(Tuple{$(types...)})) end const MACRO_KWARGS = [:launch] diff --git a/lib/KernelInterface/src/utils.jl b/lib/KernelInterface/src/utils.jl index b424a9ac0..8603aef96 100644 --- a/lib/KernelInterface/src/utils.jl +++ b/lib/KernelInterface/src/utils.jl @@ -62,3 +62,9 @@ function assign_args!(code, _args) return vars, var_exprs end + +## generated functions + +# Mark the body a generator returns for inlining. `@inline` on a generated function doesn't +# apply to a body that the generator returns with an explicit `return`, which Runic inserts. +inline_body(ex) = Expr(:block, Expr(:meta, :inline), ex) From 3d35ba288134b5573700eaaab572ba6406493d92 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 5 Nov 2026 18:07:37 +0100 Subject: [PATCH 2/2] Inline the generated functions that return their body explicitly `@inline @generated` functions whose generator returns the body with an explicit `return`, which Runic inserts, were not marked for inlining. On the POCL back-end that includes `arrayref_union` and `arrayset_union`, so that loads from arrays of isbits unions became calls returning the element through a stack slot plus a pointer to it. This is slower, and in some kernels the LLVM SPIR-V back-end then emits invalid SPIR-V for phis of these pointers (e.g. AcceleratedKernels' reductions with a `Union{Missing, Float32}` accumulator). Mark the bodies for inlining with `KernelInterface.inline_body`, here and in the host-side generated functions with the same problem. --- src/backend_launch.jl | 22 ++++++++++++---------- src/pocl/compiler/execution.jl | 12 +++++++----- src/pocl/device/array.jl | 34 +++++++++++++++++++--------------- src/pocl/device/runtime.jl | 2 +- src/pocl/nanoOpenCL.jl | 5 +++-- src/pocl/pocl.jl | 1 + test/codegen_checks.jl | 19 +++++++++++++++++++ 7 files changed, 62 insertions(+), 33 deletions(-) diff --git a/src/backend_launch.jl b/src/backend_launch.jl index 7d13402f7..d368cee0a 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -281,13 +281,15 @@ end n = fieldcount(args) converted = [Symbol(:arg, i) for i in 1:n] conversions = (:($(converted[i]) = KI.argconvert(backend, args[$i])) for i in 1:n) - return quote - converted_ctx = KI.argconvert(backend, ctx) - $(conversions...) - mask = f(ConstArgs(), $(converted...)) - return mark_constargs(mask, args), - argument_types(mask, converted_ctx, ($(converted...),)) - end + return KI.inline_body( + quote + converted_ctx = KI.argconvert(backend, ctx) + $(conversions...) + mask = f(ConstArgs(), $(converted...)) + return mark_constargs(mask, args), + argument_types(mask, converted_ctx, ($(converted...),)) + end + ) end # a kernel with varargs can have more arguments than its mask @@ -298,7 +300,7 @@ is_constarg(mask, i) = i <= length(mask) && mask[i] mark_constargs(_, args::Tuple) = args @inline @generated function mark_constargs(::Val{mask}, args::Tuple) where {mask} marked = (is_constarg(mask, i) ? :(ConstArg(args[$i])) : :(args[$i]) for i in 1:fieldcount(args)) - return :(($(marked...),)) + return KI.inline_body(:(($(marked...),))) end # `Tuple{map(Core.Typeof, (ctx, args...))...}`, with the `@Const` arguments constified @@ -308,11 +310,11 @@ end is_constarg(marked, i) ? :(Core.Typeof(constify(args[$i]))) : :(Core.Typeof(args[$i])) for i in 1:fieldcount(args) ) - return :(Tuple{Core.Typeof(ctx), $(types...)}) + return KI.inline_body(:(Tuple{Core.Typeof(ctx), $(types...)})) end # `(ctx, args...)` @inline @generated function prepend(ctx, args::Tuple) argexprs = (:(args[$i]) for i in 1:fieldcount(args)) - return :((ctx, $(argexprs...))) + return KI.inline_body(:((ctx, $(argexprs...)))) end diff --git a/src/pocl/compiler/execution.jl b/src/pocl/compiler/execution.jl index 61f12c54e..f446bd17f 100644 --- a/src/pocl/compiler/execution.jl +++ b/src/pocl/compiler/execution.jl @@ -156,7 +156,7 @@ end # stable for 32 or more elements @inline @generated function argument_types(args::Tuple) types = (:(Core.Typeof(clconvert(args[$i]))) for i in 1:fieldcount(args)) - return :(Tuple{$(types...)}) + return inline_body(:(Tuple{$(types...)})) end @@ -228,11 +228,13 @@ end call_tt, call_args = kernel_call_arguments(F, TT, fieldcount(args)) # the converted arguments only hold pointers to the arrays in `args` - return quote - GC.@preserve args begin - $cl.clcall(kernel.fun, $call_tt, ($(call_args...),); global_size, local_size, kernel.rng_state) + return inline_body( + quote + GC.@preserve args begin + $cl.clcall(kernel.fun, $call_tt, ($(call_args...),); global_size, local_size, kernel.rng_state) + end end - end + ) end diff --git a/src/pocl/device/array.jl b/src/pocl/device/array.jl index b90be88e4..13179bb39 100644 --- a/src/pocl/device/array.jl +++ b/src/pocl/device/array.jl @@ -115,15 +115,17 @@ end end end - return quote - selector_ptr = typetagdata(A, index) - selector = unsafe_load(selector_ptr) + return inline_body( + quote + selector_ptr = typetagdata(A, index) + selector = unsafe_load(selector_ptr) - align = alignment(A) - data_ptr = pointer(A, index) + align = alignment(A) + data_ptr = pointer(A, index) - return $ex - end + return $ex + end + ) end @device_function @inline function arrayset(A::CLDeviceArray{T}, x::T, index::Integer) where {T} @@ -145,16 +147,18 @@ end typs = Base.uniontypes(T) sel = findfirst(isequal(x), typs) - return quote - selector_ptr = typetagdata(A, index) - unsafe_store!(selector_ptr, $(UInt8(sel - 1))) + return inline_body( + quote + selector_ptr = typetagdata(A, index) + unsafe_store!(selector_ptr, $(UInt8(sel - 1))) - align = alignment(A) - data_ptr = pointer(A, index) + align = alignment(A) + data_ptr = pointer(A, index) - unsafe_store!(reinterpret(LLVMPtr{$x, AS}, data_ptr), x, 1, Val(align)) - return - end + unsafe_store!(reinterpret(LLVMPtr{$x, AS}, data_ptr), x, 1, Val(align)) + return + end + ) end # A load from memory that is known not to be written for the duration of the kernel. diff --git a/src/pocl/device/runtime.jl b/src/pocl/device/runtime.jl index 046cfb616..2c25de0eb 100644 --- a/src/pocl/device/runtime.jl +++ b/src/pocl/device/runtime.jl @@ -17,7 +17,7 @@ end @inline @generated function exception_field(::Val{field}) where {field} T = fieldtype(ExceptionInfo_st, field) offset = fieldoffset(ExceptionInfo_st, Base.fieldindex(ExceptionInfo_st, field)) - return :(reinterpret(LLVMPtr{$T, AS.CrossWorkgroup}, kernel_state().exception_info + $offset)) + return inline_body(:(reinterpret(LLVMPtr{$T, AS.CrossWorkgroup}, kernel_state().exception_info + $offset))) end # the record is shared by all work-items on the device, so these need device scope, which diff --git a/src/pocl/nanoOpenCL.jl b/src/pocl/nanoOpenCL.jl index 6a8d1fad0..233a01795 100644 --- a/src/pocl/nanoOpenCL.jl +++ b/src/pocl/nanoOpenCL.jl @@ -5,6 +5,7 @@ using ..POCL: platform, device, context, queue import pocl_standalone_jll using GPUToolbox: GPUToolbox, @gcsafe_ccall, cooperative_wait +using KernelInterface: inline_body using Printf @@ -1388,7 +1389,7 @@ end # one call per argument: a splat of more than 32 arguments isn't a direct call @inline @generated function set_args!(k::Kernel, args::Tuple) calls = (:(set_arg!(k, $i, args[$i])) for i in 1:fieldcount(args)) - return :($(calls...); nothing) + return inline_body(:($(calls...); nothing)) end # work sizes padded to the three dimensions OpenCL devices support @@ -1529,7 +1530,7 @@ end ).args ) - return ex + return inline_body(ex) end # the arguments are passed as a tuple, see `set_args!` diff --git a/src/pocl/pocl.jl b/src/pocl/pocl.jl index 3fcf9790a..a752e86f0 100644 --- a/src/pocl/pocl.jl +++ b/src/pocl/pocl.jl @@ -128,6 +128,7 @@ Base.Experimental.@MethodTable(method_table) import Core: LLVMPtr import GPUToolbox import UnsafeAtomics +using KernelInterface: inline_body # the size of an element in an array, which for isbits unions includes padding to their # alignment. on Julia 1.10, `Base.aligned_sizeof` isn't constant-folded in device code. diff --git a/test/codegen_checks.jl b/test/codegen_checks.jl index 8d0b9249a..6a5a732da 100644 --- a/test/codegen_checks.jl +++ b/test/codegen_checks.jl @@ -84,6 +84,11 @@ end @inbounds A[I] = 2 * A[I] end +@kernel function codegen_copy(A, B) + I = @index(Global, Linear) + @inbounds B[I] = A[I] +end + # a 64-bit primitive type that no backend shuffles natively primitive type CodegenBits64 64 end @@ -295,6 +300,20 @@ end end end + # Elements of an array of isbits unions are loaded and stored inline, rather than through + # calls that return the union and its type tag separately. + @testset "isbits unions" begin + A = Union{Missing, Float32}[1, missing, 3, 4] + B = similar(A) + @test @filecheck implicit_check_not = ["arrayref_union", "arrayset_union"] begin + @check "define spir_kernel void @{{.*}}gpu_codegen_copy" + @check "ret void" + @device_code_llvm debuginfo = :none codegen_copy(backend, 4)(A, B, ndrange = 4) + KernelAbstractions.synchronize(backend) + end + @test isequal(B, A) + end + # `@print` lowers to a single variadic printf call, not to one call per argument. @testset "print" begin @test @filecheck begin