Skip to content

Commit 86045d6

Browse files
author
Ivo Maatman
committed
Remove repeated kernel launch config calls, adapt FFTNormalOp to accomodate threads/blocks settings, add Cart GPU test
1 parent 9d26110 commit 86045d6

3 files changed

Lines changed: 114 additions & 28 deletions

File tree

‎ext/MRISubspaceReconCUDAExt/FFTNormalOp.jl‎

Lines changed: 20 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,19 @@ function MRISubspaceRecon.FFTNormalOp(Λ::CuArray{Tc}; cmaps=(1,), eltype_x=elty
2727
fftplan = plan_fft!(kL1, 1:length(img_shape))
2828
ifftplan = plan_ifft!(kL2, 1:length(img_shape))
2929

30-
A = MRISubspaceRecon._FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps)
30+
# Pre-compute kernel launch configuration
31+
Nk = size(Λ, 3)
32+
kL1_rs = reshape(kL1, :, Ncoeff)
33+
kL2_rs = reshape(kL2, :, Ncoeff)
34+
kernel = @cuda launch=false kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs)
35+
config = launch_configuration(kernel.fun)
36+
37+
threads_x = min(config.threads ÷ Ncoeff, Nk)
38+
threads_y = min(config.threads ÷ threads_x, Ncoeff)
39+
threads = (threads_x, threads_y)
40+
blocks = cld.((Nk, Ncoeff), threads)
41+
42+
A = MRISubspaceRecon._FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps, threads, blocks)
3143

3244
return LinearOperator(
3345
eltype_x,
@@ -43,7 +55,7 @@ function MRISubspaceRecon.FFTNormalOp(Λ::CuArray{Tc}; cmaps=(1,), eltype_x=elty
4355
end
4456

4557
## ##########################################################################
46-
# Kernel Calculation
58+
# Internal use
4759
#############################################################################
4860

4961
function calculate_kernel_cartesian(img_shape, trj::CuArray{<:Integer,3}, U; sample_mask=CUDA.ones(Bool, size(trj)[2:end]), verbose=false)
@@ -59,12 +71,12 @@ function calculate_kernel_cartesian(img_shape, trj::CuArray{<:Integer,3}, U; sam
5971
verbose && println("calculating Cartesian kernel on GPU...")
6072
t = @elapsed CUDA.@sync begin
6173
# Configure kernel launch
62-
kernel = @cuda launch=false _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep)
74+
kernel = @cuda launch=false kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep)
6375
config = launch_configuration(kernel.fun)
6476
threads = min(config.threads, Nsamp)
6577
blocks = cld(Nsamp, threads)
6678

67-
@cuda threads=threads blocks=blocks _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep)
79+
@cuda threads=threads blocks=blocks kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape, Ncoeff, Nsamp, Nt, Nrep)
6880
end
6981
verbose && println("time to compute kernel: t = $t s")
7082

@@ -74,7 +86,7 @@ function calculate_kernel_cartesian(img_shape, trj::CuArray{<:Integer,3}, U; sam
7486
return Λ
7587
end
7688

77-
function _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape::NTuple{N,Int}, Ncoeff, Nsamp, Nt, Nrep) where N
89+
function kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape::NTuple{N,Int}, Ncoeff, Nsamp, Nt, Nrep) where N
7890
is = (blockIdx().x - 1) * blockDim().x + threadIdx().x
7991

8092
if is <= Nsamp
@@ -91,7 +103,7 @@ function _kernel_cartesian_complex!(Λ_real, trj, U, sample_mask, img_shape::NTu
91103
for irep in 1:Nrep
92104
val += conj(U[it, ic1, irep]) * U[it, ic2, irep]
93105
end
94-
# CUDA atomics don't support complex directly, so split into real/imag
106+
# CUDA atomics does not support complex directly, so split into real/imag
95107
CUDA.@atomic Λ_real[1, ic1, ic2, k_idx...] += real(val)
96108
CUDA.@atomic Λ_real[2, ic1, ic2, k_idx...] += imag(val)
97109
end
@@ -126,9 +138,7 @@ function LinearAlgebra.mul!(x::CuArray, S::MRISubspaceRecon._FFTNormalOp, b, α,
126138
kL2_rs = reshape(S.kL2, :, S.Ncoeff)
127139
fill!(S.kL2, 0)
128140

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)
141+
@cuda threads=S.threads blocks=S.blocks kernel_mul_cartesian!(kL2_rs, S.Λ, kL1_rs, S.kmask_indcs)
132142

133143
# Inverse FFT and accumulate
134144
S.ifftplan * S.kL2
@@ -137,23 +147,7 @@ function LinearAlgebra.mul!(x::CuArray, S::MRISubspaceRecon._FFTNormalOp, b, α,
137147
return x
138148
end
139149

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)
150+
function kernel_mul_cartesian!(kL2_rs, Λ, kL1_rs, kmask_indcs)
157151
ik = (blockIdx().x - 1) * blockDim().x + threadIdx().x
158152
ic1 = (blockIdx().y - 1) * blockDim().y + threadIdx().y
159153

‎src/FFTNormalOp.jl‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ function FFTNormalOp(Λ; cmaps=(1,), num_fft_threads=round(Int, Threads.nthreads
5656
ktmp = @view kL1[CartesianIndices(img_shape), 1]
5757
fftplan = plan_fft!(ktmp; flags=FFTW.MEASURE, num_threads=num_fft_threads)
5858
ifftplan = plan_ifft!(ktmp; flags=FFTW.MEASURE, num_threads=num_fft_threads)
59-
A = _FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps)
59+
A = _FFTNormalOp(img_shape, Ncoeff, fftplan, ifftplan, Λ, kmask_indcs, kL1, kL2, cmaps, nothing, nothing)
6060

6161
return LinearOperator(
6262
eltype_x,
@@ -73,7 +73,7 @@ end
7373
## ##########################################################################
7474
# Internal use
7575
#############################################################################
76-
struct _FFTNormalOp{S,E,F,G,H,I,J,K}
76+
struct _FFTNormalOp{S,E,F,G,H,I,J,K,L,M}
7777
shape::S
7878
Ncoeff::Int
7979
fftplan::E
@@ -83,6 +83,8 @@ struct _FFTNormalOp{S,E,F,G,H,I,J,K}
8383
kL1::I
8484
kL2::J
8585
cmaps::K
86+
threads::L
87+
blocks::M
8688
end
8789

8890
function calculate_kernel_cartesian(img_shape, trj::AbstractArray{<:Integer,3}, U; sample_mask=trues(size(trj)[2:end]), verbose=false)

‎test/reconstruct_cart_trj_gpu.jl‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
using MRISubspaceRecon
2+
using ImagePhantoms
3+
using LinearAlgebra
4+
using IterativeSolvers
5+
using FFTW
6+
using Test
7+
using CUDA
8+
9+
## set parameters
10+
T = Float32
11+
Tint = Int32
12+
Nx = 128
13+
Nc = 4
14+
Nt = 40
15+
Ncoil = 9
16+
img_shape = (Nx, Nx)
17+
18+
## create test image
19+
x = zeros(Complex{T}, Nx, Nx, Nc)
20+
x[:,:,1] = transpose(shepp_logan(Nx))
21+
x[1:end÷2,:,1] .*= exp(1im * π/3)
22+
x[:,:,2] = shepp_logan(Nx)
23+
24+
## coil maps
25+
cmaps = ones(Complex{T}, Nx, Nx, Ncoil)
26+
[cmaps[i,:,2] .*= exp( 1im * π * i/Nx) for i ∈ axes(cmaps,1)]
27+
[cmaps[i,:,3] .*= exp(-1im * π * i/Nx) for i ∈ axes(cmaps,1)]
28+
[cmaps[:,i,4] .*= exp( 1im * π * i/Nx) for i ∈ axes(cmaps,2)]
29+
[cmaps[:,i,5] .*= exp(-1im * π * i/Nx) for i ∈ axes(cmaps,2)]
30+
[cmaps[i,:,6] .*= exp( 2im * π * i/Nx) for i ∈ axes(cmaps,1)]
31+
[cmaps[i,:,7] .*= exp(-2im * π * i/Nx) for i ∈ axes(cmaps,1)]
32+
[cmaps[:,i,8] .*= exp( 2im * π * i/Nx) for i ∈ axes(cmaps,2)]
33+
[cmaps[:,i,9] .*= exp(-2im * π * i/Nx) for i ∈ axes(cmaps,2)]
34+
35+
for i ∈ CartesianIndices(@view cmaps[:,:,1])
36+
cmaps[i,:] ./= norm(cmaps[i,:])
37+
end
38+
cmaps = [cmaps[:,:,ic] for ic=1:Ncoil]
39+
40+
## set up basis functions
41+
U = randn(Complex{T}, Nt, Nc)
42+
U,_,_ = svd(U)
43+
44+
## simulate data
45+
data = Array{Complex{T}}(undef, Nx, Nx, Nt, Ncoil)
46+
for icoil = 1:Ncoil
47+
Threads.@threads for i ∈ CartesianIndices(@view x[:,:,1])
48+
data[i,:,icoil] .= U * x[i,:] .* cmaps[icoil][i]
49+
end
50+
end
51+
52+
data .= ifftshift(data, (1, 2))
53+
data = reshape(data, size(data,1), size(data,2), size(data,3)*size(data,4)) # reshape prevents MKL backend error
54+
fft!(data, [1, 2])
55+
data = fftshift(data, (1, 2))
56+
data = reshape(data, Nx*Nx, Nt, Ncoil)
57+
58+
sample_mask = rand(Nx, Nx, Nt) .< 0.8 # sampling mask
59+
trj = collect(Iterators.product(1:Nx, 1:Nx, 1:Nt))
60+
kx = reshape(getindex.(trj, 1), (1, Nx*Nx, Nt))
61+
ky = reshape(getindex.(trj, 2), (1, Nx*Nx, Nt))
62+
trj = Tint.(cat(kx, ky; dims=1))
63+
sample_mask = reshape(sample_mask, Nx*Nx, Nt)
64+
65+
## load arrays onto device
66+
data = cu(data)
67+
trj = cu(trj)
68+
U = cu(U)
69+
cmaps = [cu(cmaps[i]) for i in eachindex(cmaps)]
70+
sample_mask = cu(sample_mask)
71+
72+
## calculate kernel and get normal operator
73+
AHA = FFTNormalOp((Nx,Nx), trj, U; cmaps, sample_mask)
74+
75+
## test that forward operator is symmetric
76+
Λ = zeros(Complex{T}, Nc, Nc, Nx^2)
77+
Λ[:,:,Array(AHA.prod!.A.kmask_indcs)] = Array(AHA.prod!.A.Λ) # to CPU for scalar indexing
78+
79+
Λ = reshape(Λ, Nc, Nc, Nx, Nx)
80+
for i ∈ CartesianIndices((Nx, Nx))
81+
@test Λ[:,:,i]' ≈ Λ[:,:,i]
82+
end
83+
84+
## test cg recon
85+
xbp = calculate_backprojection(data, trj, cmaps; U, sample_mask)
86+
xr = cg(AHA, vec(xbp); maxiter=20);
87+
xr = reshape(xr, Nx, Nx, Nc)
88+
89+
## test equivalence
90+
@test Array(xr) ≈ x rtol=1e-3

0 commit comments

Comments
 (0)