diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 8ef7c071d..f2dda8d47 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -463,6 +463,17 @@ Adapt.adapt_storage(::ConstAdaptor, a::Array) = Base.Experimental.Const(a) constify(arg) = adapt(ConstAdaptor(), arg) +# `constify` runs inside the kernel, where wrappers must be rebuilt without re-validating +# them: Adapt.jl's rules for these wrappers go through constructors whose error paths build +# strings, which does not compile for GPUs. Adapting only replaces the parent array, so the +# existing fields remain valid. +Adapt.adapt_structure(to::ConstAdaptor, A::Base.ReshapedArray) = + Base.ReshapedArray(adapt(to, parent(A)), size(A), A.mi) +@eval function Adapt.adapt_structure(to::ConstAdaptor, A::PermutedDimsArray{T, N, perm, iperm}) where {T, N, perm, iperm} + P = adapt(to, parent(A)) + return $(Expr(:new, :(PermutedDimsArray{eltype(P), N, perm, iperm, typeof(P)}), :P)) +end + ### # Backend hierarchy ### diff --git a/test/test.jl b/test/test.jl index d0d3e1d07..f01fb6ccc 100644 --- a/test/test.jl +++ b/test/test.jl @@ -175,6 +175,11 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @inbounds A[I] = B[I] end + @kernel function constarg2d(A, @Const(B)) + i, j = @index(Global, NTuple) + @inbounds A[i, j] = B[i, j] + end + @conditional_testset "Const" skip_tests begin let kernel = constarg(Backend(), 8, (1024,)) # this is poking at internals @@ -217,6 +222,25 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @test_skip false end end + + # wrapped arrays are rebuilt around the constified array inside the kernel + host = Float32.(reshape(1:25, 5, 5)) + dev = adapt(Backend(), host) + for (B, ref) in ( + (vec(view(dev, 1:4, 1:4)), vec(view(host, 1:4, 1:4))), + (reshape(view(dev, 1:4, 1:4), 2, 8), reshape(view(host, 1:4, 1:4), 2, 8)), + (reshape(view(dev, :, 2:3), 2, 5), reshape(view(host, :, 2:3), 2, 5)), + (PermutedDimsArray(dev, (2, 1)), PermutedDimsArray(host, (2, 1))), + ) + A = KernelAbstractions.zeros(Backend(), Float32, size(B)) + if ndims(B) == 1 + constarg(Backend(), 8)(A, B, ndrange = size(A)) + else + constarg2d(Backend(), (4, 4))(A, B, ndrange = size(A)) + end + synchronize(Backend()) + @test Array(A) == ref + end end @kernel function kernel_val!(a, ::Val{m}) where {m}