|
| 1 | +function MRISubspaceRecon.FFTNormalOp(img_shape, trj::CuArray{<:Integer,3}, U; cmaps=(1,), sample_mask=CUDA.ones(Bool, size(trj)[2:end])) |
| 2 | + Λ = calculate_kernel_cartesian(img_shape, trj, U; sample_mask) |
| 3 | + return MRISubspaceRecon.FFTNormalOp(Λ; cmaps) |
| 4 | +end |
| 5 | + |
| 6 | +# Wrapper for 4D data arrays |
| 7 | +function MRISubspaceRecon.FFTNormalOp(img_shape, trj::CuArray{<:Integer,4}, U; sample_mask=CUDA.ones(Bool, size(trj)[2:end]), kwargs...) |
| 8 | + trj = reshape(trj, size(trj, 1), :, size(trj, 4)) |
| 9 | + sample_mask = reshape(sample_mask, :, size(sample_mask, 3)) |
| 10 | + return MRISubspaceRecon.FFTNormalOp(img_shape, trj, U; sample_mask, kwargs...) |
| 11 | +end |
| 12 | + |
| 13 | +# GPU FFTNormalOp constructor from pre-computed Λ |
| 14 | +function MRISubspaceRecon.FFTNormalOp(Λ::CuArray{Tc}; cmaps=(1,), eltype_x=eltype(Λ)) where {Tc <: Complex} |
| 15 | + |
| 16 | + Ncoeff = size(Λ, 1) |
| 17 | + img_shape = size(Λ)[3:end] |
| 18 | + |
| 19 | + kL1 = CuArray{Tc}(undef, img_shape..., Ncoeff) |
| 20 | + kL2 = similar(kL1) |
| 21 | + |
| 22 | + @views kmask = (Λ[1, 1, CartesianIndices(img_shape)] .!= 0) |
| 23 | + kmask_indcs = findall(vec(kmask)) |
| 24 | + Λ = reshape(Λ, Ncoeff, Ncoeff, :) |
| 25 | + Λ = Λ[:, :, kmask_indcs] |
| 26 | + |
| 27 | + fftplan = plan_fft!(kL1, 1:length(img_shape)) |
| 28 | + ifftplan = plan_ifft!(kL2, 1:length(img_shape)) |
| 29 | + |
| 30 | + A = MRISubspaceRecon._FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps) |
| 31 | + |
| 32 | + return LinearOperator( |
| 33 | + eltype_x, |
| 34 | + prod(A.shape) * A.Ncoeff, |
| 35 | + prod(A.shape) * A.Ncoeff, |
| 36 | + true, |
| 37 | + true, |
| 38 | + (res, x, α, β) -> mul!(res, A, x, α, β), |
| 39 | + nothing, |
| 40 | + (res, x, α, β) -> mul!(res, A, x, α, β); |
| 41 | + S = typeof(similar(Λ, eltype_x, 0)) |
| 42 | + ) |
| 43 | +end |
| 44 | + |
| 45 | +## ########################################################################## |
| 46 | +# Kernel Calculation |
| 47 | +############################################################################# |
| 48 | + |
| 49 | +function calculate_kernel_cartesian(img_shape, trj::CuArray{<:Integer,3}, U; sample_mask=CUDA.ones(Bool, size(trj)[2:end]), verbose=false) |
| 50 | + Ncoeff = size(U, 2) |
| 51 | + Nrep = size(U, 3) # number of repetitions (defaults to 1) |
| 52 | + |
| 53 | + # For complex U, we need to accumulate real and imaginary parts separately via atomics |
| 54 | + Λ_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 |
| 55 | + |
| 56 | + Nsamp = size(trj, 2) |
| 57 | + Nt = size(trj, 3) |
| 58 | + |
| 59 | + verbose && println("calculating Cartesian kernel on GPU...") |
| 60 | + t = @elapsed CUDA.@sync begin |
| 61 | + # Configure kernel launch |
| 62 | + kernel = @cuda launch=false _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep) |
| 63 | + config = launch_configuration(kernel.fun) |
| 64 | + threads = min(config.threads, Nsamp) |
| 65 | + blocks = cld(Nsamp, threads) |
| 66 | + |
| 67 | + @cuda threads=threads blocks=blocks _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep) |
| 68 | + end |
| 69 | + verbose && println("time to compute kernel: t = $t s") |
| 70 | + |
| 71 | + # Reinterpret the 2-element real array as complex |
| 72 | + Λ = reinterpret(eltype(U), reshape(Λ_real, 2, Ncoeff * Ncoeff * prod(img_shape))) |
| 73 | + Λ = reshape(Λ, Ncoeff, Ncoeff, img_shape...) |
| 74 | + return Λ |
| 75 | +end |
| 76 | + |
| 77 | +function _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape::NTuple{N,Int}, Ncoeff, Nsamp, Nt, Nrep) where N |
| 78 | + is = (blockIdx().x - 1) * blockDim().x + threadIdx().x |
| 79 | + |
| 80 | + if is <= Nsamp |
| 81 | + for it in 1:Nt |
| 82 | + if sample_mask[is, it] |
| 83 | + # Compute k-space index with ifftshift incorporated |
| 84 | + k_idx = ntuple(Val(N)) do j |
| 85 | + mod1(Int(trj[j, is, it]) - img_shape[j] ÷ 2, img_shape[j]) |
| 86 | + end |
| 87 | + |
| 88 | + # Accumulate into kernel using atomic adds on real/imag parts |
| 89 | + for ic2 in 1:Ncoeff, ic1 in 1:Ncoeff |
| 90 | + val = zero(eltype(U)) |
| 91 | + for irep in 1:Nrep |
| 92 | + val += conj(U[it, ic1, irep]) * U[it, ic2, irep] |
| 93 | + end |
| 94 | + # CUDA atomics don't support complex directly, so split into real/imag |
| 95 | + CUDA.@atomic Λ_real[1, ic1, ic2, k_idx...] += real(val) |
| 96 | + CUDA.@atomic Λ_real[2, ic1, ic2, k_idx...] += imag(val) |
| 97 | + end |
| 98 | + end |
| 99 | + end |
| 100 | + end |
| 101 | + return |
| 102 | +end |
| 103 | + |
| 104 | +## ########################################################################## |
| 105 | +# mul! for GPU _FFTNormalOp |
| 106 | +############################################################################# |
| 107 | + |
| 108 | +function LinearAlgebra.mul!(x::CuArray, S::MRISubspaceRecon._FFTNormalOp, b, α, β) |
| 109 | + idx = CartesianIndices(S.shape) |
| 110 | + |
| 111 | + b = reshape(b, S.shape..., S.Ncoeff) |
| 112 | + if β == 0 |
| 113 | + x .= 0 |
| 114 | + else |
| 115 | + x .*= β |
| 116 | + end |
| 117 | + xr = reshape(x, S.shape..., S.Ncoeff) |
| 118 | + |
| 119 | + for cmap ∈ S.cmaps |
| 120 | + # Forward FFT: multiply by coil map and transform |
| 121 | + S.kL1[idx, :] .= cmap .* b |
| 122 | + S.fftplan * S.kL1 |
| 123 | + |
| 124 | + # Multiply by Toeplitz kernel |
| 125 | + kL1_rs = reshape(S.kL1, :, S.Ncoeff) |
| 126 | + kL2_rs = reshape(S.kL2, :, S.Ncoeff) |
| 127 | + fill!(S.kL2, 0) |
| 128 | + |
| 129 | + # Batched matrix multiplication at each masked k-space location |
| 130 | + # Each kmask_indcs[i] indexes a row in kL1_rs/kL2_rs, and Λ[:,:,i] is the kernel matrix |
| 131 | + _batched_kernel_mul!(kL2_rs, S.Λ, kL1_rs, S.kmask_indcs) |
| 132 | + |
| 133 | + # Inverse FFT and accumulate |
| 134 | + S.ifftplan * S.kL2 |
| 135 | + @views xr .+= α .* conj.(cmap) .* S.kL2[idx, :] |
| 136 | + end |
| 137 | + return x |
| 138 | +end |
| 139 | + |
| 140 | +# Batched kernel multiplication on GPU |
| 141 | +function _batched_kernel_mul!(kL2_rs::CuArray, Λ::CuArray, kL1_rs::CuArray, kmask_indcs) |
| 142 | + Ncoeff = size(Λ, 1) |
| 143 | + Nk = size(Λ, 3) |
| 144 | + |
| 145 | + kernel = @cuda launch=false _kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs) |
| 146 | + config = launch_configuration(kernel.fun) |
| 147 | + |
| 148 | + threads_x = min(config.threads ÷ Ncoeff, Nk) |
| 149 | + threads_y = min(config.threads ÷ threads_x, Ncoeff) |
| 150 | + threads = (threads_x, threads_y) |
| 151 | + blocks = cld.((Nk, Ncoeff), threads) |
| 152 | + |
| 153 | + @cuda threads=threads blocks=blocks _kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs) |
| 154 | +end |
| 155 | + |
| 156 | +function _kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs) |
| 157 | + ik = (blockIdx().x - 1) * blockDim().x + threadIdx().x |
| 158 | + ic1 = (blockIdx().y - 1) * blockDim().y + threadIdx().y |
| 159 | + |
| 160 | + if ik <= length(kmask_indcs) && ic1 <= size(kL2_rs, 2) |
| 161 | + ind = kmask_indcs[ik] |
| 162 | + acc = zero(eltype(kL2_rs)) |
| 163 | + |
| 164 | + @inbounds for ic2 in axes(Λ, 2) |
| 165 | + acc += Λ[ic1, ic2, ik] * kL1_rs[ind, ic2] |
| 166 | + end |
| 167 | + kL2_rs[ind, ic1] = acc |
| 168 | + end |
| 169 | + return |
| 170 | +end |
0 commit comments