Add GPU support to MRIReco.jl

.buildkite/pipeline.yml

@@ -0,0 +1,99 @@
+steps:
+  - label: "Nvidia GPUs -- MRIReco.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    command: |
+      julia --color=yes --project -e '
+      using Pkg
+      Pkg.develop([PackageSpec(path=pwd(), subdir="MRIBase")
+        , PackageSpec(path=pwd(), subdir="MRIFiles")
+        , PackageSpec(path=pwd(), subdir="MRISampling")
+        , PackageSpec(path=pwd(), subdir="MRISimulation")
+        , PackageSpec(path=pwd(), subdir="MRIOperators")
+        , PackageSpec(path=pwd(), subdir="MRICoilSensitivities")])
+      Pkg.add("TestEnv")
+      using TestEnv
+      TestEnv.activate();
+      Pkg.add("CUDA")
+      Pkg.instantiate()
+      include("test/gpu/cuda.jl")'
+    timeout_in_minutes: 30
+
+  #- label: "AMD GPUs -- MRIReco.jl"
+  #  plugins:
+  #    - JuliaCI/julia#v1:
+  #        version: "1.10"
+  #  agents:
+  #    queue: "juliagpu"
+  #    rocm: "*"
+  #    rocmgpu: "*"
+  #  command: |
+  #    julia --color=yes --project -e '
+  #    using Pkg
+  #    Pkg.develop([PackageSpec(path=pwd(), subdir="MRIBase")
+  #      , PackageSpec(path=pwd(), subdir="MRIFiles")
+  #      , PackageSpec(path=pwd(), subdir="MRISampling")
+  #      , PackageSpec(path=pwd(), subdir="MRISimulation")
+  #      , PackageSpec(path=pwd(), subdir="MRIOperators")
+  #      , PackageSpec(path=pwd(), subdir="MRICoilSensitivities")])
+  #    Pkg.add("TestEnv")
+  #    using TestEnv
+  #    TestEnv.activate();
+  #    Pkg.add("AMDGPU")
+  #    Pkg.instantiate()
+  #    include("test/gpu/rocm.jl")'
+  #  timeout_in_minutes: 30
+
+  - label: "Nvidia GPUs -- MRIOperators.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    command: |
+      julia --color=yes --project -e '
+      using Pkg
+      Pkg.develop([PackageSpec(path=pwd(), subdir="MRIBase")
+        , PackageSpec(path=pwd(), subdir="MRIFiles")
+        , PackageSpec(path=pwd(), subdir="MRISampling")
+        , PackageSpec(path=pwd(), subdir="MRISimulation")
+        , PackageSpec(path=pwd(), subdir="MRIOperators")
+        , PackageSpec(path=pwd(), subdir="MRICoilSensitivities")])
+      Pkg.add("TestEnv")
+      using TestEnv
+      TestEnv.activate("MRIOperators");
+      Pkg.add("CUDA")
+      Pkg.instantiate()
+      include("MRIOperators/test/gpu/cuda.jl")'
+    timeout_in_minutes: 30
+
+
+  - label: "AMD GPUs -- MRIOperators.jl"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.10"
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    command: |
+      julia --color=yes --project -e '
+      using Pkg
+      Pkg.develop([PackageSpec(path=pwd(), subdir="MRIBase")
+        , PackageSpec(path=pwd(), subdir="MRIFiles")
+        , PackageSpec(path=pwd(), subdir="MRISampling")
+        , PackageSpec(path=pwd(), subdir="MRISimulation")
+        , PackageSpec(path=pwd(), subdir="MRIOperators")
+        , PackageSpec(path=pwd(), subdir="MRICoilSensitivities")])
+      Pkg.add("TestEnv")
+      using TestEnv
+      TestEnv.activate("MRIOperators");
+      Pkg.add("AMDGPU")
+      Pkg.instantiate()
+      include("MRIOperators/test/gpu/rocm.jl")'
+    timeout_in_minutes: 30

.github/workflows/ci.yml

@@ -44,12 +44,12 @@ jobs:
         shell: julia --color=yes --project=. {0}
         run: |
           using Pkg
-          Pkg.develop(PackageSpec(path=pwd(), subdir="MRIBase"))
-          Pkg.develop(PackageSpec(path=pwd(), subdir="MRIFiles"))
-          Pkg.develop(PackageSpec(path=pwd(), subdir="MRISampling"))
-          Pkg.develop(PackageSpec(path=pwd(), subdir="MRISimulation"))
-          Pkg.develop(PackageSpec(path=pwd(), subdir="MRIOperators"))
-          Pkg.develop(PackageSpec(path=pwd(), subdir="MRICoilSensitivities"))
+          Pkg.develop([PackageSpec(path=pwd(), subdir="MRIBase")
+            , PackageSpec(path=pwd(), subdir="MRIFiles")
+            , PackageSpec(path=pwd(), subdir="MRISampling")
+            , PackageSpec(path=pwd(), subdir="MRISimulation")
+            , PackageSpec(path=pwd(), subdir="MRIOperators")
+            , PackageSpec(path=pwd(), subdir="MRICoilSensitivities")])
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - name: Run subpackage tests
@@ -78,12 +78,12 @@ jobs:
         shell: julia --color=yes --project=. {0}
         run: |
           using Pkg
-          Pkg.add(PackageSpec(path=pwd(), subdir="MRIBase"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="MRIFiles"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="MRISampling"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="MRISimulation"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="MRIOperators"))
-          Pkg.add(PackageSpec(path=pwd(), subdir="MRICoilSensitivities"))
+          Pkg.develop([PackageSpec(path=pwd(), subdir="MRIBase")
+            , PackageSpec(path=pwd(), subdir="MRIFiles")
+            , PackageSpec(path=pwd(), subdir="MRISampling")
+            , PackageSpec(path=pwd(), subdir="MRISimulation")
+            , PackageSpec(path=pwd(), subdir="MRIOperators")
+            , PackageSpec(path=pwd(), subdir="MRICoilSensitivities")])
       - run: |
           julia --project=docs -e '
             using Pkg

MRIBase/Project.toml

@@ -1,7 +1,7 @@
 name = "MRIBase"
 uuid = "f7771a9a-6e57-4e71-863b-6e4b6a2f17df"
 author = ["Tobias Knopp <tobias@knoppweb.de>"]
-version = "0.4.3"
+version = "0.4.4"
 
 [deps]
 AbstractNFFTs = "7f219486-4aa7-41d6-80a7-e08ef20ceed7"

MRIBase/src/Trajectories/3D/Cartesian3D.jl

@@ -23,7 +23,7 @@ function CartesianTrajectory3D(::Type{T}, numProfiles, numSamplingPerProfile
                   , kargs...) where T
   nodes = cartesian3dNodes(T, numProfiles, numSamplingPerProfile; numSlices=numSlices)
   times = readoutTimes(T, numProfiles, numSamplingPerProfile, numSlices; TE=TE, AQ=AQ)
-  return  Trajectory("Cartesian3D", nodes, times, TE, AQ, numProfiles, numSamplingPerProfile, numSlices, true, false)
+  return  Trajectory("Cartesian3D", nodes, times, T(TE), T(AQ), numProfiles, numSamplingPerProfile, numSlices, true, false)
 end
 
 function cartesian3dNodes(::Type{T}, numProfiles, numSamplingPerProfile

MRIOperators/Project.toml

@@ -1,9 +1,10 @@
 name = "MRIOperators"
 uuid = "fb1137e3-90a6-46ce-a672-6e1e53d120f2"
 author = ["Tobias Knopp <tobias@knoppweb.de>"]
-version = "0.2.1"
+version = "0.3.0"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 MRIBase = "f7771a9a-6e57-4e71-863b-6e4b6a2f17df"
@@ -17,15 +18,20 @@ LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
 LinearOperatorCollection = "a4a2c56f-fead-462a-a3ab-85921a5f2575"
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 
+[weakdeps]
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
 [compat]
+Adapt = "3, 4"
 StatsBase = "0.33, 0.34"
 Distributions = "0.25"
+GPUArrays = "8, 9, 10"
+JLArrays = "0.1"
 julia = "1.6"
 MRIBase = "0.4"
 Reexport = "1"
 LinearOperators = "2.3.3"
-LinearOperatorCollection = "1.1"
+LinearOperatorCollection = "2"
 NFFT = "0.13"
 FLoops = "0.2"
 LowRankApprox = "0.5"
@@ -34,11 +40,10 @@ Wavelets = "0.9, 0.10"
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 MRISimulation = "8988da37-ea20-4fa6-9af7-8a6f6f9a8970"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 
 [targets]
-test = ["Test", "MRISimulation"]
-
-
-
-
+test = ["Test", "MRISimulation", "JLArrays"]
 
+[extensions]
+MRIOperatorsGPUArraysExt = "GPUArrays"

MRIOperators/ext/MRIOperatorsGPUArraysExt/ExplicitOp.jl

@@ -0,0 +1,110 @@
+function MRIOperators.produ!(out::vecTc, x::vecTc, shape::NTuple{2,Int64},
+  nodes::matT, times::vecT, echoOffset::T,
+  disturbanceTerm::matTc) where {T, Tc <: Union{T, Complex{T}}, vecTc <: AbstractGPUArray{Tc}, matT <: AbstractGPUArray{T}, vecT <: AbstractGPUArray{T}, matTc <: AbstractGPUArray{Tc}}
+
+  factor = Tc(-2 * 1im * pi)
+  limit = prod(shape)
+  fill!(out, zero(Tc))
+  gpu_call(out, reshape(x, shape), shape, nodes, times, echoOffset, disturbanceTerm; elements = size(nodes, 2)) do ctx, out_, x_, shape_, nodes_, times_, echoOffset_, disturbanceTerm_
+    k = linear_index(ctx)
+    if !(1 <= k <= limit)
+      return nothing
+    end
+
+    for nx = 1:shape_[1]
+      for ny = 1:shape_[2]
+        phi = (nodes_[1, k] * (nx - shape_[1] / 2 - 1) + nodes_[2, k] * (ny - shape_[2] / 2 - 1))
+        e = exp(factor * phi - (times_[k] - echoOffset_) * disturbanceTerm_[nx, ny])
+        out_[k] += x_[nx, ny] * e
+      end
+    end
+    return nothing
+  end
+
+  return out
+end
+
+function MRIOperators.produ!(out::vecTc, x::vecTc, shape::NTuple{3,Int64},
+  nodes::matT, times::vecT, echoOffset::T,
+  disturbanceTerm::arrTc) where {T, Tc <: Union{T, Complex{T}}, vecTc <: AbstractGPUArray{Tc}, matT <: AbstractGPUArray{T}, vecT <: AbstractGPUArray{T}, arrTc <: AbstractGPUArray{Tc, 3}}
+
+  factor = Tc(-2 * 1im * pi)
+  limit = prod(shape)
+  fill!(out, zero(Tc))
+  gpu_call(out, reshape(x, shape), shape, nodes, times, echoOffset, disturbanceTerm; elements = size(nodes, 2)) do ctx, out_, x_, shape_, nodes_, times_, echoOffset_, disturbanceTerm_
+    k = linear_index(ctx)
+    if !(1 <= k <= limit)
+      return nothing
+    end
+
+    for nx = 1:shape_[1]
+      for ny = 1:shape_[2]
+        for nz = 1:shape_[3]
+          phi = (nodes_[1, k] * (nx - shape_[1] / 2 - 1) + nodes_[2, k] * (ny - shape_[2] / 2 - 1) + nodes_[3, k] * (nz - shape_[3] / 2 - 1))
+          e = exp(factor * phi - (times_[k] - echoOffset_) * disturbanceTerm_[nx, ny, nz])
+          out_[k] += x_[nx, ny, nz] * e
+        end
+      end
+    end
+    return nothing
+  end
+
+  return out
+end
+
+function MRIOperators.ctprodu!(out::vecTc, x::vecTc, shape::NTuple{2,Int64},
+  nodes::matT, times::vecT, echoOffset::T,
+  disturbanceTerm::matTc) where {T, Tc <: Union{T, Complex{T}}, vecTc <: AbstractGPUArray{Tc}, matT <: AbstractGPUArray{T}, vecT <: AbstractGPUArray{T}, matTc <: AbstractGPUArray{Tc}}
+
+  factor = Tc(-2 * 1im * pi)
+  limit = prod(shape)
+  fill!(out, zero(Tc))
+  gpu_call(reshape(out, shape), x, shape, nodes, times, echoOffset, disturbanceTerm; elements = limit) do ctx, out_, x_, shape_, nodes_, times_, echoOffset_, disturbanceTerm_
+    linearIndex = linear_index(ctx)
+    if !(1 <= linearIndex <= limit)
+      return nothing
+    end
+
+    cartIndex = CartesianIndices(shape)[linearIndex]
+
+    nx = cartIndex[1]
+    ny = cartIndex[2]
+    for k = 1:size(nodes_, 2)
+      phi = (nodes_[1,k]*(nx-shape_[1]/2-1)+nodes_[2,k]*(ny-shape_[2]/2-1))
+      e = exp(factor * phi - (times_[k]-echoOffset_)*disturbanceTerm_[cartIndex])
+      out_[cartIndex] += x_[k] * conj(e)
+    end
+    return nothing
+  end
+
+  return out
+end
+
+function MRIOperators.ctprodu!(out::vecTc, x::vecTc, shape::NTuple{3,Int64},
+  nodes::matT, times::vecT, echoOffset::T,
+  disturbanceTerm::arrTc) where {T, Tc <: Union{T, Complex{T}}, vecTc <: AbstractGPUArray{Tc}, matT <: AbstractGPUArray{T}, vecT <: AbstractGPUArray{T}, arrTc <: AbstractGPUArray{Tc, 3}}
+
+  factor = Tc(-2 * 1im * pi)
+  limit = prod(shape)
+  fill!(out, zero(Tc))
+  gpu_call(reshape(out, shape), x, shape, nodes, times, echoOffset, disturbanceTerm; elements = limit) do ctx, out_, x_, shape_, nodes_, times_, echoOffset_, disturbanceTerm_
+    linearIndex = linear_index(ctx)
+    if !(1 <= linearIndex <= limit)
+      return nothing
+    end
+
+    cartIndex = CartesianIndices(shape)[linearIndex]
+
+    nx = cartIndex[1]
+    ny = cartIndex[2]
+    nz = cartIndex[3]
+    for k = 1:size(nodes_, 2)
+      phi = (nodes_[1,k]*(nx-shape_[1]/2-1)+nodes_[2,k]*(ny-shape_[2]/2-1)+nodes_[3,k]*(nz-shape_[3]/2-1))
+      e = exp(factor * phi - (times_[k]-echoOffset_)*disturbanceTerm_[cartIndex])
+      out_[cartIndex] += x_[k] * conj(e)
+    end
+    return nothing
+  end
+
+  return out
+end

MRIOperators/ext/MRIOperatorsGPUArraysExt/FieldmapNFFTOp.jl

@@ -0,0 +1,41 @@
+function MRIOperators.produ_inner!(K, C::matT, A::matT, shape, d::Vector{vecT}, s::vecT, sp, plan, idx, x_::vecT, p::Vector{arrT}) where {T, vecT <: AbstractGPUVector{T}, matT <: AbstractGPUMatrix{T}, arrT <: AbstractGPUArray{T}}
+  for κ=1:K
+    if !isempty(idx[κ])
+      p[κ][:] .= C[κ,:] .* x_
+      mul!(d[κ], plan[κ], p[κ])  
+      # Assumption: l is unique per idx[κ]
+      gpu_call(idx[κ], d[κ], view(A, :, κ), s) do ctx, indices, d_, A_, s_
+        k = @linearidx(indices)
+        l = indices[k]
+        s_[l] += d_[k]*A_[l]
+        return nothing
+      end
+    end
+  end
+
+  return
+end
+
+
+function MRIOperators.ctprodu_inner!(K, C::matT, A::matT, shape, d::Vector{vecT}, y::vecT, sp, plan, idx, x::vecT, p::Vector{arrT}) where {T, vecT <: AbstractGPUVector{T}, matT <: AbstractGPUMatrix{T}, arrT <: AbstractGPUArray{T}}
+
+  for κ=1:K
+    if !isempty(idx[κ])
+      gpu_call(idx[κ], d[κ], view(A, :, κ), x) do ctx, indices, d_, A_, x_
+        k = @linearidx(indices)
+        l = indices[k]
+        d_[k] = conj(A_[l]) * x_[l]
+        return nothing
+      end
+      mul!(p[κ], adjoint(plan[κ]), d[κ])
+
+      gpu_call(p[κ], view(C, κ, :), y) do ctx, p_, C_, y_
+        k = @linearidx(p_)
+        y_[k] += conj(C_[k]) * p_[k]
+        return nothing
+      end
+    end
+  end
+
+  return
+end

MRIOperators/ext/MRIOperatorsGPUArraysExt/MRIOperatorsGPUArraysExt.jl

@@ -0,0 +1,13 @@
+module MRIOperatorsGPUArraysExt
+
+using MRIOperators, GPUArrays, MRIOperators.FFTW, MRIOperators.LinearAlgebra
+
+include("ExplicitOp.jl")
+include("Shutter.jl")
+include("SensitivityOp.jl")
+include("FieldmapNFFTOp.jl")
+
+MRIOperators.fftParams(::Type{<:AbstractGPUArray}) = (;)
+
+
+end # module