Move KomaCore to use KernelAbstractions.jl (to support CUDA, AMD, Metal, etc)

.github/workflows/ci.yml

@@ -44,8 +44,8 @@ jobs:
         shell: julia --color=yes --project {0}
         run: |
           using Pkg
-          Pkg.develop(PackageSpec(; path="./KomaMRICore"))
           Pkg.develop(PackageSpec(; path="./KomaMRIPlots"))
+          Pkg.develop(PackageSpec(; path="./KomaMRICore"))
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         with:

KomaMRICore/Project.toml

@@ -1,42 +1,56 @@
 name = "KomaMRICore"
 uuid = "4baa4f4d-2ae9-40db-8331-a7d1080e3f4e"
 authors = ["Carlos Castillo Passi <cncastillo@uc.cl>"]
-version = "0.7.5"
+version = "0.8.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 MRIBase = "f7771a9a-6e57-4e71-863b-6e4b6a2f17df"
 MRIFiles = "5a6f062f-bf45-497d-b654-ad17aae2a530"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Scanf = "6ef1bc8b-493b-44e1-8d40-549aa65c4b41"
 ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+KomaCUDAExt = "CUDA"
+
 [compat]
 Adapt = "3"
-CUDA = "3"
+CUDA = "4"
 FileIO = "1"
 Functors = "0.4"
 HDF5 = "0.16"
 Interpolations = "0.13, 0.14"
 JLD2 = "0.4"
+KernelAbstractions = "0.9"
 MAT = "0.10"
 MRIBase = "0.3"
 MRIFiles = "0.1"
+PackageExtensionCompat = "1"
 Parameters = "0.12"
+Preferences = "1"
 ProgressMeter = "1"
 Reexport = "1"
 Scanf = "0.5"
 ThreadsX = "0.1"
 julia = "1.6"
+
+[extras]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

KomaMRICore/ext/KomaCUDAExt.jl

@@ -0,0 +1,57 @@
+module KomaCUDAExt
+
+    using KomaMRICore
+    using KomaMRICore: KomaCPUAdaptor, KomaCUDAAdaptor, fmap
+    using CUDA
+    using Adapt
+    import KomaMRICore: _cuda       #To extend it
+    import Adapt: adapt_storage
+
+    const USE_CUDA = Ref{Union{Nothing, Bool}}(nothing)
+
+    function check_use_cuda()
+        if !isnothing(USE_CUDA[])
+            return
+        end
+        USE_CUDA[] = CUDA.functional()
+        if !USE_CUDA[]
+            @info """
+            The CUDA function is being called but CUDA.jl is not functional.
+            Defaulting back to the CPU. (No action is required if you want to run on the CPU).
+            """ maxlog=1
+        end
+        return
+    end
+
+    function __init__()
+        KomaMRICore.CUDA_LOADED[] = true
+    end
+
+    # GPU adaptor
+    adapt_storage(to::KomaCUDAAdaptor, x) = CUDA.cu(x)
+
+    function _cuda(x)
+        check_use_cuda()
+        USE_CUDA[] || return x
+        fmap(x -> Adapt.adapt(KomaCUDAAdaptor(), x), x; exclude=KomaMRICore._isleaf)
+    end
+
+    """
+    print_gpus()
+
+    Simple function to print the CUDA devices available in the host.
+    """
+    print_gpus() = begin
+        check_use_cuda()
+        if USE_CUDA[]
+            println( "$(length(devices())) CUDA capable device(s)." )
+            for (i,d) = enumerate(devices())
+                u = i == 1 ? "*" : " "
+                println( "  ($(i-1)$u) $(name(d))")
+            end
+        else
+            println("0 CUDA capable devices(s).")
+        end
+    end
+
+end

KomaMRICore/src/KomaMRICore.jl

@@ -3,14 +3,13 @@ module KomaMRICore
 #IMPORT PACKAGES
 import Base.*, Base.+, Base.-, Base./, Base.vcat, Base.size, Base.abs, Base.getproperty
 #General
-using Reexport, ThreadsX, Pkg
+using Reexport, ThreadsX
 #Printing
 using Scanf, ProgressMeter
 #Datatypes
 using Parameters
 #Simulation
 using Interpolations
-using CUDA
 #Reconstruction
 using MRIBase, MRIFiles
 @reexport using MRIBase: Profile, RawAcquisitionData, AcquisitionData, AcquisitionHeader
@@ -72,6 +71,26 @@ export get_M0, get_M1, get_M2, get_kspace
 include("sequences/PulseDesigner.jl")
 export PulseDesigner
 
+using Pkg
 __VERSION__ = VersionNumber(Pkg.TOML.parsefile(joinpath(@__DIR__, "..", "Project.toml"))["version"])
 
+using PackageExtensionCompat
+function __init__()
+    @require_extensions
+end
+#PrecompileTools
+# using PrecompileTools
+# @setup_workload begin
+#     obj = brain_phantom2D();
+#     sys = Scanner();
+#     seq = read_seq(joinpath(dirname(pathof(KomaMRICore)), 
+#         "../../examples/3.koma_paper/comparison_accuracy/sequences/EPI/epi_100x100_TE100_FOV230.seq"));
+#     simParams = KomaMRICore.default_sim_params();
+#     simParams["gpu"] = true;
+#     @compile_workload begin
+#         raw = simulate(obj, seq, sys; simParams);
+#         #abs.(raw.profiles[1].data)
+#     end
+# end
+
 end

KomaMRICore/src/simulation/BlochKA/KAtest.jl

@@ -0,0 +1,35 @@
+using KernelAbstractions
+
+# Simple kernel for matrix multiplication
+@kernel function matmul_kernel!(a, b, c)
+    i, j = @index(Global, NTuple)
+    # creating a temporary sum variable for matrix multiplication
+    tmp_sum = zero(eltype(c))
+    for k = 1:size(a)[2]
+        tmp_sum += a[i,k] * b[k, j]
+    end
+    c[i,j] = tmp_sum
+end
+
+# Creating a wrapper kernel for launching with error checks
+function matmul!(a, b, c)
+    if size(a)[2] != size(b)[1]
+        println("Matrix size mismatch!")
+        return nothing
+    end
+    backend = KernelAbstractions.get_backend(a)
+    kernel! = matmul_kernel!(backend)
+    kernel!(a, b, c, ndrange=size(c)) 
+end
+
+# Backend
+using CUDA
+# backend = CPU()
+backend = CUDA.CUDABackend()
+a = KernelAbstractions.ones( backend, Float32, 256, 123)
+b = KernelAbstractions.ones( backend, Float32, 123, 45)
+c = KernelAbstractions.zeros(backend, Float32, 10, 45)
+
+matmul!(a,b,c)
+KernelAbstractions.synchronize(backend)
+c

KomaMRICore/src/simulation/BlochKA/TestKomaExt.jl

@@ -0,0 +1,13 @@
+using KomaMRICore
+obj = brain_phantom2D();
+sys = Scanner();
+seq = read_seq(joinpath(dirname(pathof(KomaMRICore)), "../../examples/3.koma_paper/comparison_accuracy/sequences/EPI/epi_100x100_TE100_FOV230.seq"));
+simParams = KomaMRICore.default_sim_params();
+raw = simulate(obj, seq, sys; simParams);
+
+using KomaMRICore, CUDA
+obj = brain_phantom2D();
+sys = Scanner();
+seq = read_seq(joinpath(dirname(pathof(KomaMRICore)), "../../examples/3.koma_paper/comparison_accuracy/sequences/EPI/epi_100x100_TE100_FOV230.seq"));
+simParams = KomaMRICore.default_sim_params();
+raw = simulate(obj, seq, sys; simParams);