diff --git a/ext/LinearOperatorGPUArraysExt/GradientOp.jl b/ext/LinearOperatorGPUArraysExt/GradientOp.jl index ea8e79f..948472c 100644 --- a/ext/LinearOperatorGPUArraysExt/GradientOp.jl +++ b/ext/LinearOperatorGPUArraysExt/GradientOp.jl @@ -2,7 +2,7 @@ function LinearOperatorCollection.grad!(res::vecT, img::vecT, shape::NTuple{N,In res = reshape(res, shape .- Tuple(di)) backend = get_backend(res) - @kernel cpu = false inbounds = true function grad_kernel!(res, img, di) + @kernel inbounds = true function grad_kernel!(res, img, di) idx = @index(Global, Cartesian) res[idx] = img[idx] - img[idx + di] end @@ -33,12 +33,12 @@ function LinearOperatorCollection.grad_t!(res::vecT, g::vecT, shape::NTuple{N,In return res end -@kernel cpu = false inbounds = true function grad_t_kernel_1!(res, g, di) +@kernel inbounds = true function grad_t_kernel_1!(res, g, di) idx = @index(Global, Cartesian) res[idx] += g[idx] end -@kernel cpu = false inbounds = true function grad_t_kernel_2!(res, g, di) +@kernel inbounds = true function grad_t_kernel_2!(res, g, di) idx = @index(Global, Cartesian) res[idx + di] -= g[idx] end diff --git a/src/GradientOp.jl b/src/GradientOp.jl index 8752616..85c3d3c 100644 --- a/src/GradientOp.jl +++ b/src/GradientOp.jl @@ -11,12 +11,41 @@ directional gradient operator along the dimensions `dims` for an array of size ` # Optional Keyword argument * `dims` - dimension(s) along which the gradient is applied; default is `1:length(shape)` + * `S` - storage type of the operator, e.g. `Vector{T}` (default) or `CuArray{T, 1}` + +# Moving to another array backend + + Adapt.adapt(CuArray, op) + +returns a new gradient operator for the target array type (e.g. `Array`, `JLArray`, `CuArray` +or `ROCArray`), and `copy(op)` returns a new one with the same storage type. Both rebuild the +operator and its temporary storage from `shape` and the directions taken from `dims` at +construction; a later change of the vector passed as `dims`, or an iterator consumed by the +construction, has no effect. `op` itself is not modified. The directions stay on the host. +The adapted storage must have the element type of `op`: an element-typed target such as +`CuArray{Float32}` for a `Float64` operator throws an `ArgumentError`. An operator built with +an `S` of another element type cannot be adapted to any target and throws an `ArgumentError`; +`copy` keeps it as it is. """ function GradientOp(::Type{T}; shape::NTuple{N,Int}, dims=1:length(shape), kwargs...) where {T <: Number, N} return GradientOpImpl(T, shape, dims; kwargs...) end +struct GradientOpImpl{T, vecT <: AbstractVector, O <: AbstractLinearOperator{T}, D, N, W} <: GradientOp{T} + op::O + shape::NTuple{N, Int} + # `Int` for a single direction, otherwise the tuple of directions taken from `dims` at construction + dims::D + # temporary vector of the adjoint for several directions, `nothing` for a single direction + tmp::W + function GradientOpImpl(op::O, shape::NTuple{N, Int}, dims::D, tmp::W) where {T, O <: AbstractLinearOperator{T}, D, N, W} + return new{T, storage_type(op), O, D, N, W}(op, shape, dims, tmp) + end +end + function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dims; S = Vector{T}) where N + # Iterate `dims` once: replaying the operator must not depend on the caller's container later on + dims = Tuple(dims) dirs = CartesianIndex{N}[] cols = Int64[] for dim in dims @@ -34,11 +63,12 @@ function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dims; S = Vector{T}) wher tmp = S(undef, ncol) - return LinearOperator{T}(nrow, ncol, false, false, + op = LinearOperator{T}(nrow, ncol, false, false, (res,x) -> (grad!(res,x,shape,dirs, dims, dim_ends)), (res,x) -> (grad_t!(res,x,shape,dirs, dims, dim_ends, tmp)), (res,x) -> (grad_t!(res,x,shape,dirs, dims, dim_ends, tmp)), S = S) + return GradientOpImpl(op, shape, dims, tmp) end function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dim::Int; S = Vector{T}) where N @@ -48,11 +78,41 @@ function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dim::Int; S = Vector{T}) δ[dim] = 1 δ = Tuple(δ) dir = CartesianIndex(δ) - return LinearOperator{T}(nrow, ncol, false, false, + op = LinearOperator{T}(nrow, ncol, false, false, (res,x) -> (grad!(res,x,shape,dir)), (res,x) -> (grad_t!(res,x,shape,dir)), (res,x) -> (grad_t!(res,x,shape,dir)), S = S) + return GradientOpImpl(op, shape, dim, nothing) +end + +function Base.getproperty(op::GradientOpImpl, field::Symbol) + if in(field, (:op, :shape, :dims, :tmp)) + return getfield(op, field) + else + return getproperty(getfield(op, :op), field) + end +end +Base.setproperty!(op::GradientOpImpl, field::Symbol, value) = setproperty!(op.op, field, value) + +storage_type(::GradientOpImpl{T, vecT}) where {T, vecT} = vecT + +# Rebuild the operator, its closures and temporary vector from the shape and the directions +_rebuild_gradient(op::GradientOpImpl, S) = GradientOpImpl(eltype(op), op.shape, op.dims; S = S) + +Base.copy(op::GradientOpImpl) = _rebuild_gradient(op, storage_type(op)) + +function Adapt.adapt_structure(to, op::GradientOpImpl) + T = eltype(op) + # A rebuild with other storage would change the precision of the adjoint's temporary vector + eltype(storage_type(op)) == T || throw(ArgumentError( + "cannot adapt GradientOp with storage element type $(eltype(storage_type(op))) " * + "and operator element type $(T); construct it with matching element types")) + target_storage = _adapt_target_storage(to, op.Mv, "GradientOp") + eltype(target_storage) == T || throw(ArgumentError( + "unsupported Adapt target $(to) for GradientOp: storage element type $(eltype(target_storage)) " * + "differs from the operator's element type $(T)")) + return _rebuild_gradient(op, typeof(target_storage)) end function grad!(res::T, img::U, shape, dirs, dims, dim_ends) where {T<:AbstractVector, U<:AbstractVector} diff --git a/test/testAdapt.jl b/test/testAdapt.jl index 8528559..66d151f 100644 --- a/test/testAdapt.jl +++ b/test/testAdapt.jl @@ -54,6 +54,81 @@ function custom_point_transform(x::AbstractFloat) end custom_point_transform(x::Tuple) = map(custom_point_transform, x) +gradient_test_data(::Type{T}, n) where {T <: Integer} = rand(T(-5):T(5), n) +gradient_test_data(::Type{T}, n) where {T} = rand(T, n) + +gradient_matches(actual, expected, ::Type{T}) where {T <: Integer} = Array(actual) == expected +function gradient_matches(actual, expected, ::Type{T}) where {T} + tol = real(T) === Float32 ? (rtol = 5f-5, atol = 5f-6) : (rtol = 1e-11, atol = 1e-12) + return isapprox(Array(actual), expected; tol...) +end + +function gradient_inputs(op) + T = eltype(op) + m, n = size(op) + return (x = gradient_test_data(T, n), y = gradient_test_data(T, m), + x0 = gradient_test_data(T, n), y0 = gradient_test_data(T, m), α = T(2), β = T(3)) +end + +# Host copies of the forward, transpose and adjoint actions, each with 3- and 5-argument `mul!` +function gradient_actions(op, storage, inputs) + (; x, y, x0, y0, α, β) = inputs + results = Any[] + for (A, input, init) in ((op, x, y0), (transpose(op), y, x0), (adjoint(op), y, x0)) + out = storage(zero(init)) + mul!(out, A, storage(input)) + push!(results, Array(out)) + out = storage(copy(init)) + mul!(out, A, storage(input), α, β) + push!(results, Array(out)) + end + return results +end + +gradient_actions_match(op, storage, inputs, expected) = + all(gradient_matches(actual, reference, eltype(op)) + for (actual, reference) in zip(gradient_actions(op, storage, inputs), expected)) + +# Give the source scratch deterministic contents before it is snapshotted +function initialize_gradient_tmp!(op) + op.tmp === nothing || copyto!(op.tmp, gradient_test_data(eltype(op.tmp), length(op.tmp))) + return op +end + +snapshot_gradient(op) = (Mv = Array(op.Mv), Mtu = Array(op.Mtu), + tmp = op.tmp === nothing ? nothing : Array(op.tmp), + counters = (op.nprod, op.ntprod, op.nctprod)) + +function check_gradient_snapshot(op, snapshot) + @test Array(op.Mv) == snapshot.Mv + @test Array(op.Mtu) == snapshot.Mtu + if snapshot.tmp === nothing + @test op.tmp === nothing + else + @test Array(op.tmp) == snapshot.tmp + end + @test (op.nprod, op.ntprod, op.nctprod) == snapshot.counters +end + +function check_gradient_rebuild(rebuilt, source, array_type) + @test rebuilt isa LinearOperatorCollection.GradientOpImpl + @test getfield(rebuilt, :op) !== getfield(source, :op) + check_adapted_storage(rebuilt, array_type) + @test typeof(rebuilt.Mv) <: array_type + if source.tmp === nothing + @test rebuilt.tmp === nothing + else + @test typeof(rebuilt.tmp) <: array_type + @test rebuilt.tmp !== source.tmp + end + @test rebuilt.dims === source.dims + @test rebuilt.shape == source.shape + @test size(rebuilt) == size(source) + @test eltype(rebuilt) == eltype(source) +end + +storage_vector_type(array_type, T) = typeof(array_type(zeros(T, 0))) + @testset "Adapt reconstruction" begin N = 16 @@ -421,6 +496,114 @@ custom_point_transform(x::Tuple) = map(custom_point_transform, x) @test adapted.fftplan !== source_normal.fftplan end end + + @testset "GradientOp rebuilds closures and scratch" begin + cases = (((7,), nothing, (1,)), ((7,), 1, 1), + ((4, 3), nothing, (1, 2)), ((4, 3), 2, 2), ((4, 3), 1:2, (1, 2)), + ((4, 3), (2, 1), (2, 1)), ((4, 3), (2, 1, 2), (2, 1, 2)), ((4, 3), (), ()), + ((3, 1, 2), nothing, (1, 2, 3)), ((3, 1, 2), 2, 2)) + directions = array_type === Array ? ((Array, Array),) : ((Array, array_type), (array_type, Array)) + for (from, to) in directions, T in (Float32, Float64, ComplexF32, ComplexF64, Int32), + (shape, dims, expected_dims) in cases + S = storage_vector_type(from, T) + source = dims === nothing ? GradientOp(T; shape, S) : GradientOp(T; shape, dims, S) + @test source.dims === expected_dims + inputs = gradient_inputs(source) + expected = gradient_actions(source, from, inputs) + initialize_gradient_tmp!(source) + snapshot = snapshot_gradient(source) + + adapted = Adapt.adapt(to, source) + copied = copy(source) + check_gradient_rebuild(adapted, source, to) + check_gradient_rebuild(copied, source, from) + @test gradient_actions_match(adapted, to, inputs, expected) + @test gradient_actions_match(copied, from, inputs, expected) + check_gradient_snapshot(source, snapshot) + @test gradient_actions_match(source, from, inputs, expected) + end + + @test_throws BoundsError GradientOp(Float32; shape=(4, 3), dims=3) + end + + @testset "GradientOp element-typed targets" begin + directions = array_type === Array ? ((Array, Array),) : ((Array, array_type), (array_type, Array)) + for (from, to) in directions, dims in ((1, 2), 2) + source = GradientOp(Float64; shape=(4, 3), dims, S=storage_vector_type(from, Float64)) + inputs = gradient_inputs(source) + expected = gradient_actions(source, from, inputs) + + adapted = Adapt.adapt(to{Float64}, source) + check_gradient_rebuild(adapted, source, to) + @test gradient_actions_match(adapted, to, inputs, expected) + + @test_throws ArgumentError Adapt.adapt(to{Float32}, source) + @test_throws ArgumentError Adapt.adapt(to{ComplexF64}, source) + end + + # An operator whose storage element type differs from its own is rejected for any target + for dims in ((1, 2), 2) + mixed = GradientOp(Float64; shape=(4, 3), dims, S=Vector{Float32}) + inputs = gradient_inputs(mixed) + expected = gradient_actions(mixed, Array, inputs) + for to in unique((array_type, Array, array_type{Float64}, array_type{Float32}, Array{Float64})) + @test_throws ArgumentError Adapt.adapt(to, mixed) + end + copied = copy(mixed) + @test LinearOperators.storage_type(copied) == Vector{Float32} + @test all(actual == reference for (actual, reference) in zip(gradient_actions(copied, Array, inputs), expected)) + end + end + + @testset "GradientOp replays the directions taken at construction" begin + shape = (4, 3) + T = Float32 + x = rand(T, prod(shape)) + reference = GradientOp(T; shape, dims=(2, 1)) + y = rand(T, size(reference, 1)) + expected = (reference * x, adjoint(reference) * y) + acts_like_reference(op) = gradient_matches(op * array_type(x), expected[1], T) && + gradient_matches(adjoint(op) * array_type(y), expected[2], T) + + live_dims = [2, 1] + mutated = GradientOp(T; shape, dims=live_dims) + reverse!(live_dims) + @test mutated.dims === (2, 1) + @test acts_like_reference(Adapt.adapt(array_type, mutated)) + @test acts_like_reference(Adapt.adapt(array_type, copy(mutated))) + + consumed = GradientOp(T; shape, dims=Iterators.Stateful([2, 1, 2])) + @test consumed.dims === (2, 1, 2) + x3 = rand(T, prod(shape)) + expected3 = GradientOp(T; shape, dims=(2, 1, 2)) * x3 + @test gradient_matches(Adapt.adapt(array_type, consumed) * array_type(x3), expected3, T) + @test gradient_matches(Adapt.adapt(array_type, copy(consumed)) * array_type(x3), expected3, T) + + factory = createLinearOperator(GradientOp{T}; shape, dims=(2, 1)) + @test acts_like_reference(Adapt.adapt(array_type, factory)) + + weights = rand(T, size(reference, 1)) + compositions = (ProdOp(WeightingOp(weights), reference), normalOperator(reference), + DiagOp(reference, 2)) + for source in compositions + xs = rand(T, size(source, 2)) + forward = source * xs + adapted = Adapt.adapt(array_type, source) + check_adapted_storage(adapted, array_type) + @test gradient_matches(adapted * array_type(xs), forward, T) + # The normal operator defines no adjoint action; it is self-adjoint + if !(source isa LinearOperatorCollection.NormalOpImpl) + ys = rand(T, size(source, 1)) + backward = adjoint(source) * ys + @test gradient_matches(adjoint(adapted) * array_type(ys), backward, T) + end + end + @test compositions[3].ops[1].tmp !== compositions[3].ops[2].tmp + diag = Adapt.adapt(array_type, compositions[3]) + @test all(op isa LinearOperatorCollection.GradientOpImpl for op in diag.ops) + @test diag.ops[1].tmp !== diag.ops[2].tmp + @test all(typeof(op.tmp) <: array_type for op in diag.ops) + end end end