Merge pull request #9 from JuliaImageRecon/gpuOps

GPU Support for Operators

Author: nHackel

.buildkite/pipeline.yml

@@ -0,0 +1,19 @@
+steps:
+  - label: "Nvidia GPUs -- LinearOperators.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    command: |
+      julia --color=yes --project -e '
+      using Pkg
+      Pkg.add("TestEnv")
+      using TestEnv
+      TestEnv.activate();
+      Pkg.add("CUDA")
+      Pkg.add("CuNFFT")
+      Pkg.instantiate()
+      include("test/gpu/cuda.jl")'
+    timeout_in_minutes: 30

.github/workflows/Breakage.yml

@@ -0,0 +1,74 @@
+name: Breakage 
+# Based on: https://github.com/JuliaSmoothOptimizers/LinearOperators.jl/blob/main/.github/workflows/Breakage.yml
+on:
+    pull_request:
+        branches:
+            - main
+  
+jobs:
+  break:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        pkg: [
+          "JuliaImageRecon/RegularizedLeastSquares.jl",
+          "MagneticResonanceImaging/MRIReco.jl"
+        ]
+        pkgversion: [latest, stable]
+
+    steps:
+      - uses: actions/checkout@v2
+
+      # Install Julia
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: 1
+          arch: x64
+      - uses: actions/cache@v1
+        env:
+          cache-name: cache-artifacts
+        with:
+          path: ~/.julia/artifacts
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
+          restore-keys: |
+            ${{ runner.os }}-test-${{ env.cache-name }}-
+            ${{ runner.os }}-test-
+            ${{ runner.os }}-
+      - uses: julia-actions/julia-buildpkg@v1
+
+      # Breakage test
+      - name: 'Breakage of ${{ matrix.pkg }}, ${{ matrix.pkgversion }} version'
+        env:
+          URL: ${{ matrix.pkg }}
+          VERSION: ${{ matrix.pkgversion }}
+        run: |
+          set -v
+          mkdir -p ./pr
+          echo "${{ github.event.number }}" > ./pr/NR
+          git clone https://github.com/$URL
+          export PKG=$(echo $URL | cut -f2 -d/)
+          cd $PKG
+          if [ $VERSION == "stable" ]; then
+            TAG=$(git tag -l "v*" --sort=-creatordate | head -n1)
+            if [ -z "$TAG" ]; then
+              TAG="no_tag"
+            else
+              git checkout $TAG
+            fi
+          else
+            TAG=$VERSION
+          fi
+          export TAG
+          julia -e 'using Pkg;
+            PKG, TAG, VERSION = ENV["PKG"], ENV["TAG"], ENV["VERSION"]
+            joburl = joinpath(ENV["GITHUB_SERVER_URL"], ENV["GITHUB_REPOSITORY"], "actions/runs", ENV["GITHUB_RUN_ID"])
+            TAG == "no_tag" && error("Not tag for $VERSION")
+            pkg"activate .";
+            pkg"instantiate";
+            pkg"dev ../";
+            if TAG == "latest"
+              global TAG = chomp(read(`git rev-parse --short HEAD`, String))
+            end
+            pkg"build";
+            pkg"test";'

Project.toml

@@ -13,24 +13,36 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
+Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+RadonKA = "86de8297-835b-47df-b249-c04e8db91db5"
 
 [compat]
 julia = "1.9"
+GPUArrays = "8, 9, 10"
+JLArrays = "0.1"
 NFFT = "0.13"
 LinearOperators = "2.3.3"
+RadonKA = "0.6"
 Wavelets = "0.9, 0.10"
 Reexport = "1.0"
 FFTW = "1.0"
 
 [weakdeps]
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 NFFT = "efe261a4-0d2b-5849-be55-fc731d526b0d"
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+RadonKA = "86de8297-835b-47df-b249-c04e8db91db5"
 
 [targets]
-test = ["Test", "FFTW", "Wavelets", "NFFT"]
+test = ["Test", "FFTW", "Wavelets", "NFFT", "JLArrays", "RadonKA"]
 
 [extensions]
 LinearOperatorNFFTExt = ["NFFT", "FFTW"]
 LinearOperatorFFTWExt = "FFTW"
 LinearOperatorWaveletExt = "Wavelets"
+LinearOperatorGPUArraysExt = "GPUArrays"
+LinearOperatorRadonKAExt = "RadonKA"

ext/LinearOperatorFFTWExt/FFTOp.jl

@@ -1,6 +1,6 @@
 export FFTOpImpl
 
-mutable struct FFTOpImpl{T} <: FFTOp{T}
+mutable struct FFTOpImpl{T, vecT, P <: AbstractFFTs.Plan{T}, IP <: AbstractFFTs.Plan{T}} <: FFTOp{T}
   nrow :: Int
   ncol :: Int
   symmetric :: Bool
@@ -14,10 +14,10 @@ mutable struct FFTOpImpl{T} <: FFTOp{T}
   args5 :: Bool
   use_prod5! :: Bool
   allocated5 :: Bool
-  Mv5 :: Vector{T}
-  Mtu5 :: Vector{T}
-  plan
-  iplan
+  Mv5 :: vecT
+  Mtu5 :: vecT
+  plan :: P
+  iplan :: IP
   shift::Bool
   unitary::Bool
 end
@@ -34,13 +34,14 @@ returns an operator which performs an FFT on Arrays of type T
 * `shape::Tuple`  - size of the array to transform
 * (`shift=true`)  - if true, fftshifts are performed
 * (`unitary=true`)  - if true, FFT is normalized such that it is unitary
+* (`S = Vector{T}`) - type of temporary vector, change to use on GPU
+* (`kwargs...`) - keyword arguments given to fft plan
 """
-function LinearOperatorCollection.FFTOp(T::Type; shape::NTuple{D,Int64}, shift::Bool=true, unitary::Bool=true, cuda::Bool=false) where D
+function LinearOperatorCollection.FFTOp(T::Type; shape::NTuple{D,Int64}, shift::Bool=true, unitary::Bool=true, S = Array{Complex{real(T)}}, kwargs...) where D
 
-  #tmpVec = cuda ? CuArray{T}(undef,shape) : Array{Complex{real(T)}}(undef, shape)
-  tmpVec = Array{Complex{real(T)}}(undef, shape)
-  plan = plan_fft!(tmpVec; flags=FFTW.MEASURE)
-  iplan = plan_bfft!(tmpVec; flags=FFTW.MEASURE)
+  tmpVec = similar(S(undef, 0), shape...)
+  plan = plan_fft!(tmpVec; kwargs...)
+  iplan = plan_bfft!(tmpVec; kwargs...)
 
   if unitary
     facF = T(1.0/sqrt(prod(shape)))
@@ -50,39 +51,25 @@ function LinearOperatorCollection.FFTOp(T::Type; shape::NTuple{D,Int64}, shift::
     facB = T(1.0)
   end
 
-  let shape_=shape, plan_=plan, iplan_=iplan, tmpVec_=tmpVec, facF_=facF, facB_=facB
+  let shape_ = shape, plan_ = plan, iplan_ = iplan, tmpVec_ = tmpVec, facF_ = facF, facB_ = facB
 
-  if shift
-    return FFTOpImpl{T}(prod(shape), prod(shape), false, false
-              , (res, x) -> fft_multiply_shift!(res, plan_, x, shape_, facF_, tmpVec_) 
-              , nothing
-              , (res, x) -> fft_multiply_shift!(res, iplan_, x, shape_, facB_, tmpVec_) 
-              , 0, 0, 0, true, false, true, T[], T[]
-              , plan
-              , iplan
-              , shift
-              , unitary)
-  else
-    return FFTOpImpl{T}(prod(shape), prod(shape), false, false
-            , (res, x) -> fft_multiply!(res, plan_, x, facF_, tmpVec_) 
-            , nothing
-            , (res, x) -> fft_multiply!(res, iplan_, x, facB_, tmpVec_)
-            , 0, 0, 0, true, false, true, T[], T[]
-            , plan
-            , iplan
-            , shift
-            , unitary)
-  end
+    fun! = fft_multiply!
+    if shift
+      fun! = fft_multiply_shift!
+    end
+
+    return FFTOpImpl(prod(shape), prod(shape), false, false, (res, x) -> fun!(res, plan_, x, shape_, facF_, tmpVec_),
+        nothing, (res, x) -> fun!(res, iplan_, x, shape_, facB_, tmpVec_),
+        0, 0, 0, true, false, true, similar(tmpVec, 0), similar(tmpVec, 0), plan, iplan, shift, unitary)
   end
 end
 
-function fft_multiply!(res::AbstractVector{T}, plan::P, x::AbstractVector{Tr}, factor::T, tmpVec::Array{T,D}) where {T, Tr, P<:AbstractFFTs.Plan, D}
-  tmpVec[:] .= x
-  plan * tmpVec
+function fft_multiply!(res::AbstractVector{T}, plan::P, x::AbstractVector{Tr}, ::NTuple{D}, factor::T, tmpVec::AbstractArray{T,D}) where {T, Tr, P<:AbstractFFTs.Plan, D}
+  plan * copyto!(tmpVec, x)
   res .= factor .* vec(tmpVec)
 end
 
-function fft_multiply_shift!(res::AbstractVector{T}, plan::P, x::AbstractVector{Tr}, shape::NTuple{D}, factor::T, tmpVec::Array{T,D}) where {T, Tr, P<:AbstractFFTs.Plan, D}
+function fft_multiply_shift!(res::AbstractVector{T}, plan::P, x::AbstractVector{Tr}, shape::NTuple{D}, factor::T, tmpVec::AbstractArray{T,D}) where {T, Tr, P<:AbstractFFTs.Plan, D}
   ifftshift!(tmpVec, reshape(x,shape))
   plan * tmpVec
   fftshift!(reshape(res,shape), tmpVec)
@@ -91,5 +78,5 @@ end
 
 
 function Base.copy(S::FFTOpImpl)
-  return FFTOp(eltype(S); shape=size(S.plan), shift=S.shift, unitary=S.unitary)
+  return FFTOp(eltype(S); shape=size(S.plan), shift=S.shift, unitary=S.unitary, S = LinearOperators.storage_type(S)) # TODO loses kwargs...
 end

ext/LinearOperatorFFTWExt/LinearOperatorFFTWExt.jl

@@ -1,6 +1,6 @@
 module LinearOperatorFFTWExt
 
-using LinearOperatorCollection, FFTW
+using LinearOperatorCollection, FFTW, FFTW.AbstractFFTs
 
 include("FFTOp.jl")
 include("DCTOp.jl")

ext/LinearOperatorGPUArraysExt/GradientOp.jl

@@ -0,0 +1,54 @@
+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
+
+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 .- Tuple(di))
+
+  fill!(res, zero(T))
+  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
+
+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

ext/LinearOperatorGPUArraysExt/LinearOperatorGPUArraysExt.jl

@@ -0,0 +1,8 @@
+module LinearOperatorGPUArraysExt
+
+using LinearOperatorCollection, GPUArrays
+
+include("GradientOp.jl")
+
+
+end # module

ext/LinearOperatorNFFTExt/LinearOperatorNFFTExt.jl

@@ -1,6 +1,6 @@
 module LinearOperatorNFFTExt
 
-using LinearOperatorCollection, NFFT, FFTW
+using LinearOperatorCollection, NFFT, NFFT.AbstractNFFTs, FFTW, FFTW.AbstractFFTs
 
 include("NFFTOp.jl")