Improve GradOp performance on GPU for multiple dims

Author: nHackel

ext/LinearOperatorGPUArraysExt/GradientOp.jl

@@ -1,38 +1,54 @@
-function LinearOperatorCollection.grad!(res::vecT, img::vecT, shape, dim) where {vecT <: AbstractGPUVector}
-  δ = zeros(Int, length(shape))
-  δ[dim] = 1
-  δ = Tuple(δ)
-  di = CartesianIndex(δ)
-
-  gpu_call(reshape(res, shape .- δ), reshape(img,shape), di) do ctx, res_, img_, di_
-    idx = @cartesianidx(res_)
-    @inbounds res_[idx] = img_[idx] - img_[idx + di_]
-    return nothing  
+function LinearOperatorCollection.grad!(res::vecT, img::vecT, shape::NTuple{N,Int64}, di::CartesianIndex{N}) where {vecT <: AbstractGPUVector, N}
+  res = reshape(res, shape .- Tuple(di))
+
+  if length(res) > 0
+    gpu_call(grad_kernel!, res, reshape(img,shape), di)
   end
-  
+
   return res
 end
 
-# adjoint of directional gradients
-function LinearOperatorCollection.grad_t!(res::vecT, g::vecT, shape::NTuple{N,Int64}, dim::Int64) where {T, vecT <: AbstractGPUVector{T}, N}
-  δ = zeros(Int, length(shape))
-  δ[dim] = 1
-  δ = Tuple(δ)
-  di = CartesianIndex(δ)
+function grad_kernel!(ctx, res, img, di)
+  idx = @cartesianidx(res)
+  @inbounds res[idx] = img[idx] - img[idx + di]
+  return nothing  
+end
 
+# adjoint of directional gradients
+function LinearOperatorCollection.grad_t!(res::vecT, g::vecT, shape::NTuple{N,Int64}, di::CartesianIndex{N}) where {T, vecT <: AbstractGPUVector{T}, N}
   res_ = reshape(res,shape)
-  g_ = reshape(g, shape .- δ)
+  g_ = reshape(g, shape .- Tuple(di))
 
   fill!(res, zero(T))
-  gpu_call(res_, g_, di, elements = length(g)) do ctx, res_k, g_k, di_k
-    idx = @cartesianidx(g_k)
-    @inbounds res_k[idx]  = g_k[idx]
-    return nothing  
+  if length(g_) > 0
+    gpu_call(grad_t_kernel_1!, res_, g_, di, elements = length(g))
+    gpu_call(grad_t_kernel_2!, res_, g_, di, elements = length(g))
   end
+end
 
-  gpu_call(res_, g_, di, elements = length(g)) do ctx, res_k, g_k, di_k
-    idx = @cartesianidx(g_k)
-    @inbounds res_k[idx + di_k] -= g_k[idx]
-    return nothing  
-  end
+function grad_t_kernel_1!(ctx, res, g, di)
+  idx = @cartesianidx(g)
+  @inbounds res[idx] += g[idx]
+  return nothing  
+end
+
+function grad_t_kernel_2!(ctx, res, g, di)
+  idx = @cartesianidx(g)
+  @inbounds res[idx + di] -= g[idx]
+  return nothing  
 end
+
+function LinearOperatorCollection.grad_t!(res::vecT, g::vecT, shape::NTuple{N,Int64}, dirs, dims, dim_ends, tmp) where {T, vecT <: AbstractGPUVector{T}, N}
+  dim_start = 1
+  res = reshape(res, shape)
+
+  fill!(res, zero(eltype(res)))
+  for (i, di) in enumerate(dirs)
+    g_ = reshape(view(g, dim_start:dim_ends[i]), shape .- Tuple(di))
+    if length(g_) > 0
+      gpu_call(grad_t_kernel_1!, res, g_, di, elements = length(g))
+      gpu_call(grad_t_kernel_2!, res, g_, di, elements = length(g))
+    end  
+    dim_start = dim_ends[i] + 1
+  end
+end

src/GradientOp.jl

@@ -16,46 +16,80 @@ function GradientOp(::Type{T}; shape::NTuple{N,Int}, dims=1:length(shape), kwarg
   return GradientOpImpl(T, shape, dims; kwargs...)
 end
 
-function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dims; kwargs...) where N
-  return vcat([GradientOpImpl(T, shape, dim; kwargs...) for dim ∈ dims]...)
+function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dims; S = Vector{T}) where N
+  dirs = CartesianIndex{N}[]
+  cols = Int64[]
+  for dim in dims
+    δ = zeros(Int32, N)
+    δ[dim] = 1
+    δ = NTuple{N}(δ)
+    di = CartesianIndex(δ)
+    push!(dirs, di)
+    push!(cols, div((shape[dim]-1)*prod(shape), shape[dim]))
+  end
+  dim_ends = accumulate(+, cols)
+  
+  nrow = sum(cols)
+  ncol = prod(shape)
+
+  tmp = S(undef, ncol)
+  
+  return 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)
 end
 
 function GradientOpImpl(T::Type, shape::NTuple{N,Int}, dim::Int; S = Vector{T}) where N
   nrow = div( (shape[dim]-1)*prod(shape), shape[dim] )
   ncol = prod(shape)
+  δ = zeros(Int, length(shape))
+  δ[dim] = 1
+  δ = Tuple(δ)
+  dir = CartesianIndex(δ)
   return LinearOperator{T}(nrow, ncol, false, false,
-                          (res,x) -> (grad!(res,x,shape,dim)),
-                          (res,x) -> (grad_t!(res,x,shape,dim)),
-                          (res,x) -> (grad_t!(res,x,shape,dim)),
+                          (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)
 end
 
+function grad!(res::T, img::U, shape, dirs, dims, dim_ends) where {T<:AbstractVector, U<:AbstractVector}
+  dim_start = 1
+
+  for (i, dir) in enumerate(dirs)
+    grad!(view(res, dim_start:dim_ends[i]), img, shape, dir)
+    dim_start = dim_ends[i] + 1
+  end
+end
+
 # directional gradients
-function grad!(res::T, img::U, shape, dim) where {T<:AbstractVector, U<:AbstractVector}
+function grad!(res::T, img::U, shape::NTuple{N,Int64}, di::CartesianIndex{N}) where {N, T<:AbstractVector, U<:AbstractVector}
   img_ = reshape(img,shape)
 
-  δ = zeros(Int, length(shape))
-  δ[dim] = 1
-  δ = Tuple(δ)
-  di = CartesianIndex(δ)
-
-  res_ = reshape(res, shape .- δ)
+  res_ = reshape(res, shape .- Tuple(di))
 
   Threads.@threads for i ∈ CartesianIndices(res_)
     @inbounds res_[i] = img_[i] - img_[i + di]
   end
 end
 
+function grad_t!(res::T, g::U, shape, dirs, dims, dims_end, tmp) where {T<:AbstractVector, U<:AbstractVector}
+  dim_start = 1
 
-# adjoint of directional gradients
-function grad_t!(res::T, g::U, shape::NTuple{N,Int64}, dim::Int64) where {T<:AbstractVector, U<:AbstractVector, N}
-  δ = zeros(Int, length(shape))
-  δ[dim] = 1
-  δ = Tuple(δ)
-  di = CartesianIndex(δ)
+  fill!(res, zero(eltype(res)))
+  for (i, dir) in enumerate(dirs)
+    grad_t!(tmp, view(g, dim_start:dims_end[i]), shape, dir)
+    dim_start = dims_end[i] + 1
+    res .= res .+ tmp 
+  end
+end
 
+# adjoint of directional gradients
+function grad_t!(res::T, g::U, shape::NTuple{N,Int64}, di::CartesianIndex{N}) where {N, T<:AbstractVector, U<:AbstractVector}
   res_ = reshape(res,shape)
-  g_ = reshape(g, shape .- δ)
+  g_ = reshape(g, shape .- Tuple(di))
 
   res_ .= 0
   Threads.@threads for i ∈ CartesianIndices(g_)