Merge branch 'master' into delete_LowRankApprox

MRIOperators/Project.toml

@@ -19,15 +19,21 @@ Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 
 [weakdeps]
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 
 [extensions]
 MRIOperatorsGPUArraysExt = "GPUArrays"
 
 [compat]
 Adapt = "3, 4"
-FLoops = "0.2"
-GPUArrays = "8, 9, 10"
-JLArrays = "0.1"
+StatsBase = "0.33, 0.34"
+GPUArrays = "11"
+JLArrays = "0.2"
+KernelAbstractions = "0.9"
+julia = "1.6"
+MRIBase = "0.4"
+Reexport = "1"
+LinearOperators = "2.3"
 LinearOperatorCollection = "2"
 LinearOperators = "2.3.3"
 MRIBase = "0.4"
@@ -44,3 +50,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "MRISimulation", "JLArrays"]
+
+[extensions]
+MRIOperatorsGPUArraysExt = ["GPUArrays", "KernelAbstractions"]

MRIOperators/ext/MRIOperatorsGPUArraysExt/ExplicitOp.jl

@@ -5,21 +5,18 @@ function MRIOperators.produ!(out::vecTc, x::vecTc, shape::NTuple{2,Int64},
   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
+  @kernel inbounds = true cpu = false function produ_kernel!(out, x, shape, nodes, times, echoOffset, disturbanceTerm)
+    k = @index(Global, Linear)
+    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
+  kernel! = produ_kernel!(get_backend(out))
+  kernel!(out, reshape(x, shape), shape, nodes, times, echoOffset, disturbanceTerm; ndrange = limit)
 
   return out
 end
@@ -31,23 +28,20 @@ function MRIOperators.produ!(out::vecTc, x::vecTc, shape::NTuple{3,Int64},
   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
+  @kernel inbounds = true cpu = false function produ_kernel!(out, x, shape, nodes, times, echoOffset, disturbanceTerm)
+    k = @index(Global, Linear)
+    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
+  kernel! = produ_kernel!(get_backend(out))
+  kernel!(out, reshape(x, shape), shape, nodes, times, echoOffset, disturbanceTerm; ndrange = limit)
 
   return out
 end
@@ -59,23 +53,21 @@ function MRIOperators.ctprodu!(out::vecTc, x::vecTc, shape::NTuple{2,Int64},
   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
+  @kernel inbounds = true cpu = false function ctprodu_kernel!(out, x, shape, nodes, times, echoOffset, disturbanceTerm)
+    linearIndex = @index(Global, Linear)
 
     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)
+    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
+  kernel! = ctprodu_kernel!(get_backend(out))
+  kernel!(reshape(out, shape), x, shape, nodes, times, echoOffset, disturbanceTerm; ndrange = limit)
 
   return out
 end
@@ -87,24 +79,22 @@ function MRIOperators.ctprodu!(out::vecTc, x::vecTc, shape::NTuple{3,Int64},
   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
+  @kernel inbounds = true cpu = false function ctprodu_kernel!(out, x, shape, nodes, times, echoOffset, disturbanceTerm)
+    linearIndex = @index(Global, Linear)
 
     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)
+    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
+  kernel! = ctprodu_kernel!(get_backend(out))
+  kernel!(reshape(out, shape), x, shape, nodes, times, echoOffset, disturbanceTerm; ndrange = limit)
 
   return out
 end

MRIOperators/ext/MRIOperatorsGPUArraysExt/FieldmapNFFTOp.jl

@@ -1,15 +1,18 @@
 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}}
+
+  @kernel inbounds = true cpu = false function produ_inner_kernel!(indices, d, A, s)
+    k = @index(Global, Linear)
+    # Assumption: l is unique per idx[κ]
+    l = indices[k]
+    s[l] += d[k]*A[l]
+  end
+  kernel! = produ_inner_kernel!(get_backend(C))
+
   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
+      mul!(d[κ], plan[κ], p[κ])
+      kernel!(idx[κ], d[κ], view(A, :, κ), s; ndrange = length(idx[κ]))
     end
   end
 
@@ -19,21 +22,26 @@ 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}}
 
+
+  @kernel inbounds = true cpu = false function ctprodu_inner_1!(indices, d, A, x)
+    k = @index(Global, Linear)
+    l = indices[k]
+    d[k] = conj(A[l]) * x[l]
+  end
+  kernel1! = ctprodu_inner_1!(get_backend(C))
+
+  @kernel inbounds = true cpu = false function ctprodu_inner_2!(p, C, y)
+    k = @index(Global, Linear)
+    y[k] += conj(C[k]) * p[k]
+  end
+  kernel2! = ctprodu_inner_2!(get_backend(C))
+
+
   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
+      kernel1!(idx[κ], d[κ], view(A, :, κ), x; ndrange = length(idx[κ]))
       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
+      kernel2!(p[κ], view(C, κ, :), y; ndrange = length(p[κ]))
     end
   end
 

MRIOperators/ext/MRIOperatorsGPUArraysExt/MRIOperatorsGPUArraysExt.jl

@@ -1,6 +1,7 @@
 module MRIOperatorsGPUArraysExt
 
-using MRIOperators, GPUArrays, MRIOperators.FFTW, MRIOperators.LinearAlgebra
+using MRIOperators, GPUArrays, KernelAbstractions
+using MRIOperators.FFTW, MRIOperators.LinearAlgebra
 
 include("ExplicitOp.jl")
 include("Shutter.jl")

MRIOperators/ext/MRIOperatorsGPUArraysExt/SensitivityOp.jl

@@ -4,11 +4,12 @@ function MRIOperators.prod_smap!(y::vecT, smaps::matT, x::vecT, numVox, numChan,
 
   @assert size(smaps) == (size(y_,1), size(y_,2))
 
-  gpu_call(y_, x_, smaps) do ctx, ygpu, xgpu, smapsgpu
-    cart = @cartesianidx(ygpu)
-    ygpu[cart] = xgpu[cart[1],cart[3]] * smapsgpu[cart[1],cart[2]]
-    return nothing
+  @kernel inbounds = true cpu = false function smap_kernel!(y, x, smaps)
+    cart = @index(Global, Cartesian)
+    y[cart] = x[cart[1], cart[3]] * smaps[cart[1], cart[2]]
   end
+  kernel! = smap_kernel!(get_backend(y))
+  kernel!(y_, x_, smaps; ndrange = size(y_))
   return y
 end
 
@@ -20,12 +21,13 @@ function MRIOperators.ctprod_smap!(y::vecT, smapsC::matT, x::vecT, numVox, numCh
 
   y_ .= 0
   # Inner loop to avoid race condition
-  gpu_call(y_, x_, smapsC, numChan) do ctx, ygpu, xgpu, smapsCgpu, limit
-    cart = @cartesianidx(ygpu)
+  @kernel inbounds = true cpu = false function smap_kernel_adjoint!(y, x, smapsC, limit)
+    cart = @index(Global, Cartesian)
     for i = 1:limit
-      ygpu[cart] += xgpu[cart[1], i, cart[2]] * smapsCgpu[cart[1],i]
+      y[cart] += x[cart[1], i, cart[2]] * smapsC[cart[1],i]
     end
-    return nothing
   end
+  kernel! = smap_kernel_adjoint!(get_backend(y))
+  kernel!(y_, x_, smapsC, numChan; ndrange = size(y_))
   return y
 end

MRIOperators/ext/MRIOperatorsGPUArraysExt/Shutter.jl

@@ -1,31 +1,27 @@
+@kernel inbounds = true cpu = false function circularShutterKernel(I, center, radius)
+  cart = @index(Global, Cartesian)
+  if sqrt(sum((Tuple(cart) .-center).^2)) > radius
+    I[cart] = 0
+  end
+end
+
 function MRIOperators.circularShutter!(I::AbstractGPUArray, radiusFactor::Number=1.0)
   imgSize = size(I)
   center = imgSize./2.0
   radius = maximum(center) * radiusFactor
-  # Applying filtering
-  gpu_call(I, center, radius) do ctx, I_, center_, radius_
-    cart = @cartesianidx(I_)
-    if sqrt(sum((Tuple(cart) .-center_).^2)) > radius_
-        I[cart] = 0
-    end
-    return nothing
-  end
+  kernel! = circularShutterKernel(get_backend(I))
+  kernel!(I, center, radius; ndrange = imgSize)
   return I
 end
 
 
+
 function MRIOperators.circularShutterFreq!(I::AbstractGPUArray, radiusFactor::T=1.0) where T<:Number
   imgSize = size(I)
   fftI = fftshift(fft(I))
   center = imgSize./2.0
   radius = maximum(center) * radiusFactor
-  # Applying filtering
-  gpu_call(fftI, center, radius) do ctx, fftI_, center_, radius_
-    cart = @cartesianidx(fftI_)
-    if sqrt(sum((Tuple(cart) .-center_).^2)) > radius_
-      fftI_[cart] = 0
-    end
-    return nothing
-  end
+  kernel! = circularShutterKernel(get_backend(I))
+  kernel!(fftI, center, radius; ndrange = size(fftI))
   return MRIOperators.preserveType(I, ifft(ifftshift(fftI)))
 end

Project.toml

@@ -23,7 +23,7 @@ GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 AxisArrays = "0.4.6"
 FFTW = "1.0"
 FLoops = "0.2"
-GPUArrays = "8, 9, 10"
+GPUArrays = "8, 9, 10, 11"
 ImageUtils = "0.2.8"
 MRIBase = "0.3, 0.4"
 MRIOperators = "0.3"
@@ -32,7 +32,7 @@ ProgressMeter = "1.2"
 Reexport = "0.2, 1"
 RegularizedLeastSquares = "0.16"
 Unitful = "1.2"
-julia = "1.6"
+julia = "1.10"
 
 [extras]
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"