Merge pull request #79 from MagneticResonanceImaging/CartFFTNormalOp

Cartesian GPU code

Author: IvoMaatman

README.md

@@ -9,7 +9,7 @@
 
 MRISubspaceRecon.jl is a Julia package that implements sub-space reconstructions as described by [Jon Tamir et al.](https://onlinelibrary.wiley.com/doi/abs/10.1002/mrm.26102) and by us in the paper [Low-rank alternating direction method of multipliers reconstruction for MR fingerprinting](https://doi.org/10.1002/mrm.26639) as *low rank inversion*.
 
-Currently, the package supports Cartesian and non-Cartesian trajectories on CPU, and non-Cartesian ones for NVIDIA GPUs. The package is still work in progress. The ultimate goal of this package is to provide a comprehensive Julia implementation of MRI subspace reconstructions, similar to our [Matlab package](https://bitbucket.org/asslaender/nyu_mrf_recon).
+Currently, the package supports Cartesian and non-Cartesian trajectories on single- or multithreaded CPUs, as well as NVIDIA GPUs. The package is still under active development. Its ultimate goal is to provide a comprehensive Julia implementation for MRI subspace reconstructions, analagous to our [Matlab package](https://bitbucket.org/asslaender/nyu_mrf_recon).
 
 
 [docs-img]: https://img.shields.io/badge/docs-latest%20release-blue.svg

docs/src/index.md

@@ -4,7 +4,7 @@ CurrentModule = MRISubspaceRecon
 
 # MRISubspaceRecon.jl
 
-[MRISubspaceRecon.jl](https://github.com/MagneticResonanceImaging/MRISubspaceRecon.jl) package package aims to enable rapid iterative reconstruction
+The [MRISubspaceRecon.jl](https://github.com/MagneticResonanceImaging/MRISubspaceRecon.jl) package aims to enable rapid iterative reconstruction 
 of Cartesian and non-Cartesian MRI data using subspace modeling [1,2]. Particular care is given to enable the reconstruction of large numbers of subspace coefficients along with large image grid sizes.
 
 For compatibility with other Julia packages, such as [IterativeSolvers.jl](https://github.com/JuliaLinearAlgebra/IterativeSolvers.jl) and [RegularizedLeastSquares.jl](https://github.com/JuliaImageRecon/RegularizedLeastSquares.jl), operations are defined in terms of linear operators and their effects on the data vectors. [MRISubspaceRecon.jl](https://github.com/MagneticResonanceImaging/MRISubspaceRecon.jl) is designed to compute these objects with multi-threaded CPUs and on NVIDIA GPUs. The package further contains functions to perform GRAPPA operator gridding (GROG) [3] or to generate radial trajectories. However, all methods that require an explicit k-space trajectory as input generalize to arbitrary trajectories.
@@ -31,7 +31,7 @@ Reconstructions of non-Cartesian MRI in [MRISubspaceRecon.jl](https://github.com
 using MRISubspaceRecon
 using CUDA
 ```
-We recommend using the GPU code which can be faster by a factor of 10--20 than CPU multi-threading (for typical solvers like conjugate gradient or FISTA [7]). However, a specific GPU implementation for Cartesian MRI is still under development. In this case, one can use the CPU implementation or the non-Cartesian methods.
+We recommend using the GPU code which can be faster by a factor of 10--20 than CPU multi-threading (for typical solvers like conjugate gradient or FISTA [7]).
 
 # References
 1. Assländer J, et al. “Low rank alternating direction method of multipliers reconstruction for MR fingerprinting”. Magn Reson Med 79.1 (2018), pp. 83–96. https://doi.org/10.1002/mrm.26639

ext/MRISubspaceReconCUDAExt/BackProjection.jl

@@ -81,6 +81,98 @@ function MRISubspaceRecon.calculate_backprojection(data::CuArray{Tc,4}, trj::CuA
 end
 
 
+## ##########################################################################
+# Cartesian backprojection (Integer trajectory) on GPU
+#############################################################################
+function MRISubspaceRecon.calculate_backprojection(data::CuArray{Tc}, trj::CuArray{<:Integer,3}, cmaps::AbstractVector{<:CuArray}; U=cu(I(size(trj)[end])), sample_mask=CUDA.ones(Bool, size(trj)[2:end])) where {Tc <: Complex}
+    Ncoeff = size(U, 2)
+    img_shape = size(cmaps[1])
+    img_idx = CartesianIndices(img_shape)
+    T = real(eltype(data))
+
+    dataU_real = CUDA.zeros(T, 2, img_shape...)
+    dataU = CuArray{eltype(data)}(undef, img_shape)
+    xbp = CUDA.zeros(eltype(data), img_shape..., Ncoeff)
+
+    ifftplan = plan_ifft!(dataU)
+
+    Nsamp = size(trj, 2)
+    Nt = size(trj, 3)
+    Nrep = size(U, 3)
+
+    # Configure kernel launch
+    kernel = @cuda launch=false _cartesian_backprojection_kernel!(dataU_real, data, trj, U, sample_mask, img_shape, 1, 1, Nsamp, Nt, Nrep)
+    config = launch_configuration(kernel.fun)
+    threads = min(config.threads, Nsamp)
+    blocks = cld(Nsamp, threads)
+
+    for icoef in axes(U, 2)
+        for icoil in axes(data, 3)
+            fill!(dataU_real, 0)
+            @cuda threads=threads blocks=blocks _cartesian_backprojection_kernel!(dataU_real, data, trj, U, sample_mask, img_shape, icoef, icoil, Nsamp, Nt, Nrep)
+            dataU .= reshape(reinterpret(eltype(data), vec(dataU_real)), img_shape)
+            ifftplan * dataU
+            @views xbp[img_idx, icoef] .+= conj.(cmaps[icoil]) .* fftshift(dataU)
+        end
+    end
+    return xbp
+end
+
+function MRISubspaceRecon.calculate_backprojection(data::CuArray{Tc}, trj::CuArray{<:Integer,3}, img_shape; U=cu(I(size(trj)[end])), sample_mask=CUDA.ones(Bool, size(trj)[2:end])) where {Tc <: Complex}
+    Ncoeff = size(U, 2)
+    Ncoil = size(data, 3)
+    img_idx = CartesianIndices(img_shape)
+    T = real(eltype(data))
+
+    dataU_real = CUDA.zeros(T, 2, img_shape...)
+    dataU = CuArray{eltype(data)}(undef, img_shape)
+    xbp = CUDA.zeros(eltype(data), img_shape..., Ncoeff, Ncoil)
+
+    ifftplan = plan_ifft!(dataU)
+
+    Nsamp = size(trj, 2)
+    Nt = size(trj, 3)
+    Nrep = size(U, 3)
+
+    # Configure kernel launch
+    kernel = @cuda launch=false _cartesian_backprojection_kernel!(dataU_real, data, trj, U, sample_mask, img_shape, 1, 1, Nsamp, Nt, Nrep)
+    config = launch_configuration(kernel.fun)
+    threads = min(config.threads, Nsamp)
+    blocks = cld(Nsamp, threads)
+
+    for icoef in axes(U, 2)
+        for icoil in axes(data, 3)
+            fill!(dataU_real, 0)
+            @cuda threads=threads blocks=blocks _cartesian_backprojection_kernel!(dataU_real, data, trj, U, sample_mask, img_shape, icoef, icoil, Nsamp, Nt, Nrep)
+            dataU .= reshape(reinterpret(eltype(data), vec(dataU_real)), img_shape)
+            ifftplan * dataU
+            @views xbp[img_idx, icoef, icoil] .= fftshift(dataU)
+        end
+    end
+    return xbp
+end
+
+function _cartesian_backprojection_kernel!(dataU_real, data, trj, U, sample_mask, img_shape::NTuple{N,Int}, icoef, icoil, Nsamp, Nt, Nrep) where N
+    is = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+
+    if is <= Nsamp
+        for it in 1:Nt
+            if sample_mask[is, it]
+                k_idx = ntuple(Val(N)) do j
+                    mod1(Int(trj[j, is, it]) - img_shape[j] ÷ 2, img_shape[j])
+                end
+
+                for irep in 1:Nrep
+                    val = data[is, it, icoil, irep] * conj(U[it, icoef, irep])
+                    CUDA.@atomic dataU_real[1, k_idx...] += real(val)
+                    CUDA.@atomic dataU_real[2, k_idx...] += imag(val)
+                end
+            end
+        end
+    end
+    return
+end
+
 ## ##########################################################################
 # Internal helper functions
 #############################################################################

ext/MRISubspaceReconCUDAExt/CoilMaps.jl

@@ -48,10 +48,72 @@ function MRISubspaceRecon.calculate_coil_maps(
     return calculate_coil_maps(data, trj, img_shape; sample_mask, kwargs...)
 end
 
+## ##########################################################################
+# Cartesian coil maps (Integer trajectory) on GPU
+#############################################################################
+function MRISubspaceRecon.calculate_coil_maps(
+    data::CuArray{Complex{T}},
+    trj::CuArray{<:Integer},
+    img_shape::NTuple{N,Int};
+    U=CUDA.ones(Complex{T}, size(data, 2)),
+    sample_mask=CUDA.ones(Bool, size(trj)[2:end]),
+    kernel_size=ntuple(_ -> 6, N),
+    calib_size=img_shape .÷ (img_shape[1] ÷ 32),
+    eigThresh_1=0.01,
+    eigThresh_2=0.9,
+    nmaps=1,
+    Niter_cg=5,
+    verbose=false) where {N,T}
+
+    lower_bound = @. Int(ceil((img_shape - calib_size) / 2))
+    upper_bound = @. lower_bound + calib_size + 1
+    mask_calib = dropdims(all(Array(trj) .> lower_bound; dims=1) .& all(Array(trj) .< upper_bound; dims=1); dims=1)
+    mask_calib .&= Array(sample_mask)
+    mask_calib = cu(mask_calib)
+
+    x = MRISubspaceRecon.reconstruct_coilwise(data, trj, calib_size; U, sample_mask=mask_calib, Niter_cg)
+
+    imdims = ntuple(i -> i, length(img_shape))
+    kbp = fftshift(x, imdims)
+    fft!(kbp, imdims)
+    kbp = fftshift(kbp, imdims)
+
+    t = @elapsed begin
+        cmaps = espirit(Array(kbp), img_shape, kernel_size; eigThresh_1, eigThresh_2, nmaps)
+    end
+    verbose && println("espirit: $t s")
+
+    cmaps = [cu(cmaps[CartesianIndices(img_shape), ic, in]) for ic in axes(cmaps, length(img_shape) + 1), in = (nmaps == 1 ? 1 : 1:nmaps)]
+    return cmaps
+end
+
 
 ## ##########################################################################
 # Internal helper functions
 #############################################################################
+function MRISubspaceRecon.reconstruct_coilwise(
+    data::CuArray{Tc},
+    trj::CuArray{<:Integer,3},
+    img_shape;
+    U=CUDA.ones(Tc, size(data, 2)),
+    sample_mask=CUDA.ones(Bool, size(trj)[2:end]),
+    Niter_cg=5,
+    verbose=false) where {Tc <: Complex}
+
+    AᴴA = MRISubspaceRecon.FFTNormalOp(img_shape, trj, U[:, 1]; sample_mask)
+    xbp = MRISubspaceRecon.calculate_backprojection(data, trj, img_shape; U=U[:, 1], sample_mask)
+
+    Ncoil = size(data, 3)
+    x = CUDA.zeros(Tc, img_shape..., Ncoil)
+
+    for icoil in axes(xbp, length(img_shape) + 2)
+        bi = vec(@view xbp[CartesianIndices(img_shape), 1, icoil])
+        xi = vec(@view x[CartesianIndices(img_shape), icoil])
+        cg!(xi, AᴴA, bi; maxiter=Niter_cg, verbose)
+    end
+    return x
+end
+
 function MRISubspaceRecon.reconstruct_coilwise(
     data::CuArray{Tc,3},
     trj::CuArray{T,3},

ext/MRISubspaceReconCUDAExt/FFTNormalOp.jl

@@ -0,0 +1,164 @@
+function MRISubspaceRecon.FFTNormalOp(img_shape, trj::CuArray{<:Integer,3}, U; cmaps=(1,), sample_mask=CUDA.ones(Bool, size(trj)[2:end]))
+    Λ = calculate_kernel_cartesian(img_shape, trj, U; sample_mask)
+    return MRISubspaceRecon.FFTNormalOp(Λ; cmaps)
+end
+
+# Wrapper for 4D data arrays
+function MRISubspaceRecon.FFTNormalOp(img_shape, trj::CuArray{<:Integer,4}, U; sample_mask=CUDA.ones(Bool, size(trj)[2:end]), kwargs...)
+    trj = reshape(trj, size(trj, 1), :, size(trj, 4))
+    sample_mask = reshape(sample_mask, :, size(sample_mask, 3))
+    return MRISubspaceRecon.FFTNormalOp(img_shape, trj, U; sample_mask, kwargs...)
+end
+
+# GPU FFTNormalOp constructor from pre-computed Λ
+function MRISubspaceRecon.FFTNormalOp(Λ::CuArray{Tc}; cmaps=(1,), eltype_x=eltype(Λ)) where {Tc <: Complex}
+
+    Ncoeff = size(Λ, 1)
+    img_shape = size(Λ)[3:end]
+
+    kL1 = CuArray{Tc}(undef, img_shape..., Ncoeff)
+    kL2 = similar(kL1)
+
+    @views kmask = (Λ[1, 1, CartesianIndices(img_shape)] .!= 0)
+    kmask_indcs = findall(vec(kmask))
+    Λ = reshape(Λ, Ncoeff, Ncoeff, :)
+    Λ = Λ[:, :, kmask_indcs]
+
+    fftplan  = plan_fft!(kL1, 1:length(img_shape))
+    ifftplan = plan_ifft!(kL2, 1:length(img_shape))
+
+    # Pre-compute kernel launch configuration
+    Nk = size(Λ, 3)
+    kL1_rs = reshape(kL1, :, Ncoeff)
+    kL2_rs = reshape(kL2, :, Ncoeff)
+    kernel = @cuda launch=false kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs)
+    config = launch_configuration(kernel.fun)
+
+    threads_x = min(config.threads ÷ Ncoeff, Nk)
+    threads_y = min(config.threads ÷ threads_x, Ncoeff)
+    threads = (threads_x, threads_y)
+    blocks = cld.((Nk, Ncoeff), threads)
+
+    A = MRISubspaceRecon._FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps, threads, blocks)
+
+    return LinearOperator(
+        eltype_x,
+        prod(A.shape) * A.Ncoeff,
+        prod(A.shape) * A.Ncoeff,
+        true,
+        true,
+        (res, x, α, β) -> mul!(res, A, x, α, β),
+        nothing,
+        (res, x, α, β) -> mul!(res, A, x, α, β);
+        S = typeof(similar(Λ, eltype_x, 0))
+    )
+end
+
+## ##########################################################################
+# Internal use
+#############################################################################
+
+function calculate_kernel_cartesian(img_shape, trj::CuArray{<:Integer,3}, U; sample_mask=CUDA.ones(Bool, size(trj)[2:end]), verbose=false)
+    Ncoeff = size(U, 2)
+    Nrep = size(U, 3) # number of repetitions (defaults to 1)
+
+    # For complex U, we need to accumulate real and imaginary parts separately via atomics        
+    Λ_real = CUDA.zeros(real(eltype(U)), 2, Ncoeff, Ncoeff, img_shape...) # Use a real-valued array with 2x the leading dimension to handle complex atomics
+
+    Nsamp = size(trj, 2)
+    Nt = size(trj, 3)
+
+    verbose && println("calculating Cartesian kernel on GPU...")
+    t = @elapsed CUDA.@sync begin
+        # Configure kernel launch
+        kernel = @cuda launch=false kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep)
+        config = launch_configuration(kernel.fun)
+        threads = min(config.threads, Nsamp)
+        blocks = cld(Nsamp, threads)
+
+        @cuda threads=threads blocks=blocks kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep)
+    end
+    verbose && println("time to compute kernel: t = $t s")
+
+    # Reinterpret the 2-element real array as complex
+    Λ = reinterpret(eltype(U), reshape(Λ_real, 2, Ncoeff * Ncoeff * prod(img_shape)))
+    Λ = reshape(Λ, Ncoeff, Ncoeff, img_shape...)
+    return Λ
+end
+
+function kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape::NTuple{N,Int}, Ncoeff, Nsamp, Nt, Nrep) where N
+    is = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+
+    if is <= Nsamp
+        for it in 1:Nt
+            if sample_mask[is, it]
+                # Compute k-space index with ifftshift incorporated
+                k_idx = ntuple(Val(N)) do j
+                    mod1(Int(trj[j, is, it]) - img_shape[j] ÷ 2, img_shape[j])
+                end
+
+                # Accumulate into kernel using atomic adds on real/imag parts
+                for ic2 in 1:Ncoeff, ic1 in 1:Ncoeff
+                    val = zero(eltype(U))
+                    for irep in 1:Nrep
+                        val += conj(U[it, ic1, irep]) * U[it, ic2, irep]
+                    end
+                    # CUDA atomics does not support complex directly, so split into real/imag
+                    CUDA.@atomic Λ_real[1, ic1, ic2, k_idx...] += real(val)
+                    CUDA.@atomic Λ_real[2, ic1, ic2, k_idx...] += imag(val)
+                end
+            end
+        end
+    end
+    return
+end
+
+## ##########################################################################
+# mul! for GPU _FFTNormalOp
+#############################################################################
+
+function LinearAlgebra.mul!(x::CuArray, S::MRISubspaceRecon._FFTNormalOp, b, α, β)
+    idx = CartesianIndices(S.shape)
+
+    b = reshape(b, S.shape..., S.Ncoeff)
+    if β == 0
+        x .= 0
+    else
+        x .*= β
+    end
+    xr = reshape(x, S.shape..., S.Ncoeff)
+
+    for cmap ∈ S.cmaps
+        # Forward FFT: multiply by coil map and transform
+        S.kL1[idx, :] .= cmap .* b
+        S.fftplan * S.kL1
+
+        # Multiply by Toeplitz kernel
+        kL1_rs = reshape(S.kL1, :, S.Ncoeff)
+        kL2_rs = reshape(S.kL2, :, S.Ncoeff)
+        fill!(S.kL2, 0)
+
+        @cuda threads=S.threads blocks=S.blocks kernel_mul_cartesian!(kL2_rs, S.Λ, kL1_rs, S.kmask_indcs)
+
+        # Inverse FFT and accumulate
+        S.ifftplan * S.kL2
+        @views xr .+= α .* conj.(cmap) .* S.kL2[idx, :]
+    end
+    return x
+end
+
+function kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs)
+    ik  = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+    ic1 = (blockIdx().y - 1) * blockDim().y + threadIdx().y
+
+    if ik <= length(kmask_indcs) && ic1 <= size(kL2_rs, 2)
+        ind = kmask_indcs[ik]
+        acc = zero(eltype(kL2_rs))
+
+        @inbounds for ic2 in axes(Λ, 2)
+            acc += Λ[ic1, ic2, ik] * kL1_rs[ind, ic2]
+        end
+        kL2_rs[ind, ic1] = acc
+    end
+    return
+end

ext/MRISubspaceReconCUDAExt/MRISubspaceReconCUDAExt.jl

@@ -13,5 +13,6 @@ using CUDA
 include("NFFTNormalOp.jl")
 include("BackProjection.jl")
 include("CoilMaps.jl")
+include("FFTNormalOp.jl")
 
 end # module

src/FFTNormalOp.jl

@@ -56,7 +56,7 @@ function FFTNormalOp(Λ; cmaps=(1,), num_fft_threads=round(Int, Threads.nthreads
     ktmp = @view kL1[CartesianIndices(img_shape), 1]
     fftplan = plan_fft!(ktmp; flags=FFTW.MEASURE, num_threads=num_fft_threads)
     ifftplan = plan_ifft!(ktmp; flags=FFTW.MEASURE, num_threads=num_fft_threads)
-    A = _FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps)
+    A = _FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps, nothing, nothing)
 
     return LinearOperator(
         eltype_x,
@@ -73,7 +73,7 @@ end
 ## ##########################################################################
 # Internal use
 #############################################################################
-struct _FFTNormalOp{S,E,F,G,H,I,J,K}
+struct _FFTNormalOp{S,E,F,G,H,I,J,K,L,M}
     shape::S
     Ncoeff::Int
     fftplan::E
@@ -83,6 +83,8 @@ struct _FFTNormalOp{S,E,F,G,H,I,J,K}
     kL1::I
     kL2::J
     cmaps::K
+    threads::L
+    blocks::M
 end
 
 function calculate_kernel_cartesian(img_shape, trj::AbstractArray{<:Integer,3}, U; sample_mask=trues(size(trj)[2:end]), verbose=false)