Add multi-GPU parallelisation

Project.toml

@@ -1,36 +1,53 @@
 name = "MRIReco"
 uuid = "bdf86e05-2d2b-5731-a332-f3fe1f9e047f"
-authors = ["Tobias Knopp <tobias@knoppweb.de>"]
 version = "0.9.2"
+authors = ["Tobias Knopp <tobias@knoppweb.de>"]
 
 [deps]
+AbstractImageReconstruction = "a4b4fdbf-6459-4ec9-990d-77e1fa24a91b"
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 MRIBase = "f7771a9a-6e57-4e71-863b-6e4b6a2f17df"
 MRIOperators = "fb1137e3-90a6-46ce-a672-6e1e53d120f2"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 RegularizedLeastSquares = "1e9c538a-f78c-5de5-8ffb-0b6dbe892d23"
+RelocatableFolders = "05181044-ff0b-4ac5-8273-598c1e38db00"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+
+[extensions]
+MRIRecoGPUArraysExt = "GPUArrays"
+MRIRecoAMDGPUExt = "AMDGPU"
+MRIRecoCUDAExt = "CUDA"
+MRIRecoMetalExt = "Metal"
+
 
 [compat]
+AbstractImageReconstruction = "0.6.5"
+AMDGPU = "0.9, 1, 2"
 AxisArrays = "0.4.6"
+CUDA = "3, 4, 5"
 FFTW = "1.0"
 GPUArrays = "8, 9, 10, 11"
-OhMyThreads = "0.8"
 ImageUtils = "0.2.8"
+Metal = "1"
 MRIBase = "0.3, 0.4"
 MRIOperators = "0.3, 0.4"
+OhMyThreads = "0.8"
 PrecompileTools = "1"
 ProgressMeter = "1.2"
 Reexport = "0.2, 1"
 RegularizedLeastSquares = "0.16"
+RelocatableFolders = "1"
 Unitful = "1.2"
 julia = "1.10"
 
@@ -44,6 +61,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "ImageUtils", "MRISimulation", "MRISampling", "Scratch", "FFTW"]
-
-[extensions]
-MRIRecoGPUArraysExt = "GPUArrays"

ext/MRIRecoAMDGPUExt/MRIRecoAMDGPUExt.jl

@@ -0,0 +1,7 @@
+module MRIRecoAMDGPUExt
+
+using MRIReco, AMDGPU
+
+MRIReco.set_device!(::Type{<:ROCArray}, index) = AMDGPU.device!(index) 
+
+end

ext/MRIRecoCUDAExt/MRIRecoCUDAExt.jl

@@ -0,0 +1,7 @@
+module MRIRecoCUDAExt
+
+using MRIReco, CUDA
+
+MRIReco.set_device!(::Type{<:CuArray}, index) = CUDA.device!(index) 
+
+end

reco_plans/direct.toml

@@ -0,0 +1,6 @@
+_module = "MRIReco"
+_type = "RecoPlan{DirectMRIReco}"
+
+[parameter]
+_module = "MRIReco"
+_type = "RecoPlan{DirectMRIRecoParameter}"

reco_plans/multiCoil.toml

@@ -0,0 +1,34 @@
+_module = "MRIReco"
+_type = "RecoPlan{MultiCoilReconstruction}"
+
+[parameter]
+_module = "MRIReco"
+_type = "RecoPlan{ThreadedIterativeMRIRecoContextParameter}"
+
+    [parameter.parameter]
+    _module = "MRIReco"
+    _type = "RecoPlan{MultiCoilIterativeParameters}"
+
+        [parameter.parameter.weightingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{DensityWeightingParameters}"
+
+        [parameter.parameter.encodingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{EncodingParameters}"
+
+        [parameter.parameter.coilParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{CoilParameters}"
+
+        [parameter.parameter.solverParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{LeastSquaresSolverParameter}"
+
+            [parameter.parameter.solverParams.regParams]
+            _module = "MRIReco"
+            _type = "RecoPlan{RegularizationParameters}"
+
+                [parameter.parameter.solverParams.regParams.sparsityParams]
+                _module = "MRIReco"
+                _type = "RecoPlan{SimpleSparsityParameters}"

reco_plans/multiCoilMultiEcho.toml

@@ -0,0 +1,34 @@
+_module = "MRIReco"
+_type = "RecoPlan{MultiCoilMultiEchoReconstruction}"
+
+[parameter]
+_module = "MRIReco"
+_type = "RecoPlan{ThreadedIterativeMRIRecoContextParameter}"
+
+    [parameter.parameter]
+    _module = "MRIReco"
+    _type = "RecoPlan{MultiCoilIterativeParameters}"
+
+        [parameter.parameter.weightingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{DensityWeightingParameters}"
+
+        [parameter.parameter.encodingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{EncodingParameters}"
+
+        [parameter.parameter.coilParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{CoilParameters}"
+
+        [parameter.parameter.solverParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{LeastSquaresSolverParameter}"
+
+            [parameter.parameter.solverParams.regParams]
+            _module = "MRIReco"
+            _type = "RecoPlan{RegularizationParameters}"
+
+                [parameter.parameter.solverParams.regParams.sparsityParams]
+                _module = "MRIReco"
+                _type = "RecoPlan{SimpleSparsityParameters}"

reco_plans/multiCoilMultiEchoSubspace.toml

@@ -0,0 +1,38 @@
+_module = "MRIReco"
+_type = "RecoPlan{MultiCoilMultiEchoSubspaceReconstruction}"
+
+[parameter]
+_module = "MRIReco"
+_type = "RecoPlan{ThreadedIterativeMRIRecoContextParameter}"
+
+    [parameter.parameter]
+    _module = "MRIReco"
+    _type = "RecoPlan{MultiCoilIterativeParameters}"
+
+        [parameter.parameter.weightingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{DensityWeightingParameters}"
+
+        [parameter.parameter.encodingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{SubspaceEncodingParameters}"
+
+            [parameter.parameter.encodingParams.inner]
+            _module = "MRIReco"
+            _type = "RecoPlan{EncodingParameters}"
+
+        [parameter.parameter.coilParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{CoilParameters}"
+
+        [parameter.parameter.solverParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{LeastSquaresSolverParameter}"
+
+            [parameter.parameter.solverParams.regParams]
+            _module = "MRIReco"
+            _type = "RecoPlan{RegularizationParameters}"
+
+                [parameter.parameter.solverParams.regParams.sparsityParams]
+                _module = "MRIReco"
+                _type = "RecoPlan{SimpleSparsityParameters}"

reco_plans/multiEcho.toml

@@ -0,0 +1,30 @@
+_module = "MRIReco"
+_type = "RecoPlan{MultiEchoReconstruction}"
+
+[parameter]
+_module = "MRIReco"
+_type = "RecoPlan{ThreadedIterativeMRIRecoContextParameter}"
+
+    [parameter.parameter]
+    _module = "MRIReco"
+    _type = "RecoPlan{MultiEchoIterativeParameters}"
+
+        [parameter.parameter.weightingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{DensityWeightingParameters}"
+
+        [parameter.parameter.encodingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{EncodingParameters}"
+
+        [parameter.parameter.solverParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{LeastSquaresSolverParameter}"
+
+            [parameter.parameter.solverParams.regParams]
+            _module = "MRIReco"
+            _type = "RecoPlan{RegularizationParameters}"
+
+                [parameter.parameter.solverParams.regParams.sparsityParams]
+                _module = "MRIReco"
+                _type = "RecoPlan{SimpleSparsityParameters}"

reco_plans/standard.toml

@@ -0,0 +1,30 @@
+_module = "MRIReco"
+_type = "RecoPlan{StandardReconstruction}"
+
+[parameter]
+_module = "MRIReco"
+_type = "RecoPlan{ThreadedIterativeMRIRecoContextParameter}"
+
+    [parameter.parameter]
+    _module = "MRIReco"
+    _type = "RecoPlan{StandardIterativeParameters}"
+
+        [parameter.parameter.weightingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{DensityWeightingParameters}"
+
+        [parameter.parameter.encodingParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{EncodingParameters}"
+
+        [parameter.parameter.solverParams]
+        _module = "MRIReco"
+        _type = "RecoPlan{LeastSquaresSolverParameter}"
+
+            [parameter.parameter.solverParams.regParams]
+            _module = "MRIReco"
+            _type = "RecoPlan{RegularizationParameters}"
+
+                [parameter.parameter.solverParams.regParams.sparsityParams]
+                _module = "MRIReco"
+                _type = "RecoPlan{SimpleSparsityParameters}"

src/Algorithms/Algorithms.jl

@@ -0,0 +1,21 @@
+export AbstractMRIRecoAlgorithm
+abstract type AbstractMRIRecoAlgorithm <: AbstractImageReconstructionAlgorithm end
+
+export AbstractDirectMRIRecoAlgorithm
+abstract type AbstractDirectMRIRecoAlgorithm <: AbstractMRIRecoAlgorithm end
+
+export AbstractIterativeMRIRecoAlgorithm
+abstract type AbstractIterativeMRIRecoAlgorithm <: AbstractMRIRecoAlgorithm end
+
+export AbstractMRIRecoParameters
+abstract type AbstractMRIRecoParameters <: AbstractImageReconstructionParameters end
+
+export MRIRecoStyle
+struct MRIRecoStyle <: CustomPlanStyle end
+
+include("MRIRecoContext.jl")
+include("Storage.jl")
+include("Direct.jl")
+
+include("Parameters/Parameters.jl")
+include("Iterative/Iterative.jl")

src/Algorithms/Direct.jl

@@ -0,0 +1,94 @@
+export AbstractDirectMRIRecoParameters
+
+"""
+    AbstractDirectMRIRecoParameters <: AbstractMRIRecoParameters
+
+Abstract type for direct (non-iterative) MRI reconstruction parameters.
+Performs a direct Fourier-based image reconstruction.
+"""
+abstract type AbstractDirectMRIRecoParameters <: AbstractMRIRecoParameters end
+
+export DirectMRIReco
+@reconstruction mutable struct DirectMRIReco{P <: AbstractDirectMRIRecoParameters} <: AbstractDirectMRIRecoAlgorithm
+  @parameter parameter::P
+end
+
+export DirectMRIRecoParameter
+
+"""
+    DirectMRIRecoParameter{S, C, W, arrT} <: AbstractDirectMRIRecoParameters
+
+Performs a direct Fourier-based image reconstruction of AcquisitionData.
+
+# Arguments
+- `reconSize::S` - Size of image to reconstruct (defaults to encoding size if nothing)
+- `cmap::C` - Fieldmap for correction of off-resonance effects (optional)
+- `weights::W` - Sampling density of the trajectories in acqData (optional)
+- `arrayType::Type{arrT}` - Array type to use for reconstruction (defaults to Array)
+"""
+@parameter struct DirectMRIRecoParameter{
+    S <: Union{Nothing, NTuple{D, Int64} where D},   # User supplied or default recon size
+    C <: Union{Nothing, <:AbstractArray{<:Complex}}, # User supplied or no correction map
+    W <: Union{Nothing, <:AbstractArray{<:Complex}}, # User supplied or sampling density
+    arrT                                             # Array type
+  } <: AbstractDirectMRIRecoParameters
+  reconSize::S = nothing
+  cmap::C = nothing
+  weights::W = nothing
+  arrayType::Type{arrT} = Array
+end
+
+
+# Setup recon size
+(params::DirectMRIRecoParameter{Nothing})(algo::AbstractDirectMRIRecoAlgorithm, acqData::AcquisitionData) = params(algo, acqData, encodingSize(acqData))
+(params::DirectMRIRecoParameter)(algo::AbstractDirectMRIRecoAlgorithm, acqData::AcquisitionData) = params(algo, acqData, params.reconSize)
+
+# Prepare image
+function (params::AbstractDirectMRIRecoParameters)(algo::AbstractDirectMRIRecoAlgorithm, acqData::AcquisitionData{T}, reconSize::NTuple{D, Int64}) where {T, D}
+  encDims = ndims(trajectory(acqData))
+  if encDims!=D
+    error("reco-dimensionality $D and encoding-dimensionality $(encDims) do not match")
+  end
+
+  numContr, numChan, numSl, numRep = numContrasts(acqData), numChannels(acqData), numSlices(acqData), numRepetitions(acqData)
+  Ireco = zeros(Complex{T}, prod(reconSize), numSl, numContr, numChan, numRep)
+
+  Ireco = with(MRIRECO_CONTEXT => MRIRecoContext(reconSize, acqData, params.arrayType)) do
+    return params(algo, Ireco, acqData)
+  end
+
+  Ireco = reshape(Ireco, volumeSize(reconSize, numSl)..., numContr, numChan, numRep)
+  return makeAxisArray(Ireco, acqData)
+end
+
+# Reconstruction loop
+function (params::DirectMRIRecoParameter)(::Type{<:AbstractDirectMRIRecoAlgorithm}, Ireco::Array{Complex{T}, 5}, acqData::AcquisitionData{T}) where T
+  numSl = size(Ireco, 2)
+  numContr = size(Ireco, 3)
+  numChan = size(Ireco, 4)
+  numRep = size(Ireco, 5)
+
+  correctionMap = isnothing(params.cmap) ? similar(Ireco, 0) : params.cmap
+  reconSize = ctx_reconSize()
+  weights = isnothing(params.weights) ? samplingDensity(acqData, reconSize) : params.weights
+  arrayType = ctx_arrayType()
+  S = ctx_storageType()
+
+  for i = 1:numSl
+    F = encodingOps_simple(acqData, reconSize, slice=i, correctionMap=correctionMap, S = S)
+    for k = 1:numContr
+      for j = 1:numChan
+        for l = 1:numRep
+            kdata = arrayType(kData(acqData,k,j,i,rep=l)) .* arrayType((weights[k].^2))
+            I = adjoint(F[k]) * kdata
+            if isCircular( trajectory(acqData, k) )
+              circularShutter!(reshape(I, reconSize), 1.0)
+            end
+            Ireco[:,i,k,j,l] = Array(I)
+        end
+      end
+    end
+  end
+
+  return Ireco
+end