Inline the generated functions that return their body explicitly

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

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]

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)

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

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
 
 

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.

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

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!`

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.

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