Fix GPU signal sampling during excitaiton (KomaMRICore v0.11) (#757)

* Fix GPU reduction for partial spin groups

* Fix GPU CI test environment activation

* Skip excitation signal reduction without ADC

* Skip GPU signal reduction without ADC

* Restore Julia 1.10 GPU CI activation guard

Author: cncastillo

AGENTS.md

@@ -93,6 +93,7 @@ Pkg.test()  # KomaMRI umbrella package
 - `KomaMRICore` defaults to CPU via `KomaMRICore/test/Project.toml`. To change backend, set `[preferences.KomaMRICore].test_backend` to `CPU`, `CUDA`, `AMDGPU`, `Metal`, or `oneAPI`, or use `Pkg.test("KomaMRICore"; test_args=\`CUDA\`)`.
 - To control CPU thread count for `KomaMRICore`, use `Pkg.test("KomaMRICore"; julia_args=\`--threads=4\`)`.
 - GPU backend packages must already be available locally before running non-CPU `KomaMRICore` tests. `oneAPI` is experimental.
+- Do not commit backend-specific GPU packages such as `CUDA`, `AMDGPU`, `Metal`, or `oneAPI` to shared test projects like `KomaMRICore/test/Project.toml`. Add them only in the backend-specific CI/local test command or a temporary environment.
 - For small exploratory tests, use a temporary environment. For longer isolated tests, use a temp folder with an activated environment. Reuse the same REPL.
 - If you add behavior, add or update tests in the owning package's `runtests.jl`.
 - Use `@testitem` tags correctly:

KomaMRICore/src/simulation/SimMethods/Bloch/gpu/BlochGPU.jl

@@ -37,21 +37,24 @@ function run_spin_precession!(
 ) where {T<:Real, SM<:BlochLikeSimMethods}
     #Motion
     x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t')
+    has_adc = length(sig) > 0
 
     #Precession
     precession_kernel!(backend, groupsize)(
         pre.sig_output,
         M.xy, M.z,
         x, y, z, pre.ΔBz, p.T1, p.T2, p.ρ, UInt32(length(M.xy)),
         seq.Gx, seq.Gy, seq.Gz, seq.Δt, seq.ADC, UInt32(length(seq.t)),
-        Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)),
+        Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)), Val(has_adc),
         sim_method,
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
     )
 
     #Signal
-    AK.reduce(+, view(pre.sig_output,:,1:length(sig)); init=zero(Complex{T}), dims=1, temp=view(pre.sig_output_final,:,1:length(sig)))
-    sig .= transpose(view(pre.sig_output_final,:,1:length(sig)))
+    if has_adc
+        AK.reduce(+, view(pre.sig_output,:,1:length(sig)); init=zero(Complex{T}), dims=1, temp=view(pre.sig_output_final,:,1:length(sig)))
+        sig .= transpose(view(pre.sig_output_final,:,1:length(sig)))
+    end
 
     #Reset Spin-State (Magnetization). Only for FlowPath
     outflow_spin_reset!(M, seq.t', p.motion; replace_by=p.ρ)
@@ -71,21 +74,24 @@ function run_spin_excitation!(
 ) where {T<:Real, SM<:BlochLikeSimMethods}
     #Motion
     x, y, z = get_spin_coords(p.motion, p.x, p.y, p.z, seq.t')
+    has_adc = length(sig) > 0
 
     #Excitation
     excitation_kernel!(backend, groupsize)(
         pre.sig_output,
         M.xy, M.z,
         x, y, z, pre.ΔBz, p.T1, p.T2, p.ρ, UInt32(length(M.xy)),
         seq.Gx, seq.Gy, seq.Gz, seq.Δt, seq.Δf, seq.B1, seq.ψ, seq.ADC, UInt32(length(seq.t)),
-        Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)),
+        Val(!(p.motion isa NoMotion)), Val(supports_warp_reduction(backend)), Val(has_adc),
         sim_method,
         ndrange=(cld(length(M.xy), groupsize) * groupsize)
     )
 
     #Signal
-    AK.reduce(+, view(pre.sig_output,:,1:length(sig)); init=zero(Complex{T}), dims=1, temp=view(pre.sig_output_final,:,1:length(sig)))
-    sig .= transpose(view(pre.sig_output_final,:,1:length(sig)))
+    if has_adc
+        AK.reduce(+, view(pre.sig_output,:,1:length(sig)); init=zero(Complex{T}), dims=1, temp=view(pre.sig_output_final,:,1:length(sig)))
+        sig .= transpose(view(pre.sig_output_final,:,1:length(sig)))
+    end
 
     #Reset Spin-State (Magnetization). Only for FlowPath
     outflow_spin_reset!(M,  seq.t', p.motion; replace_by=p.ρ) # TODO: reset state inside kernel

KomaMRICore/src/simulation/SimMethods/Bloch/gpu/ExcitationKernel.jl

@@ -5,21 +5,34 @@
     M_xy::AbstractVector{Complex{T}}, M_z, 
     @Const(p_x), @Const(p_y), @Const(p_z), @Const(p_ΔBz), @Const(p_T1), @Const(p_T2), @Const(p_ρ), N_spins,
     @Const(s_Gx), @Const(s_Gy), @Const(s_Gz), @Const(s_Δt), @Const(s_Δf), @Const(s_B1), @Const(s_ψ), @Const(s_ADC), s_length,
-    ::Val{MOTION}, ::Val{USE_WARP_REDUCTION},
+    ::Val{MOTION}, ::Val{USE_WARP_REDUCTION}, ::Val{HAS_ADC},
     sim_method::SM
-) where {T, MOTION, USE_WARP_REDUCTION, SM <: BlochLikeSimMethods}
+) where {T, MOTION, USE_WARP_REDUCTION, HAS_ADC, SM <: BlochLikeSimMethods}
 
     @uniform N = @groupsize()[1]
     i_l = @index(Local, Linear)
     i_g = @index(Group, Linear)
     i = (i_g - 1u32) * UInt32(N) + i_l
 
-    sig_group_r = @localmem T USE_WARP_REDUCTION ? 32 : N
-    sig_group_i = @localmem T USE_WARP_REDUCTION ? 32 : N
-    sig_r = zero(T)
-    sig_i = zero(T)
+    sig_group_r = @localmem T HAS_ADC ? (USE_WARP_REDUCTION ? 32 : N) : 1
+    sig_group_i = @localmem T HAS_ADC ? (USE_WARP_REDUCTION ? 32 : N) : 1
 
-    if i <= N_spins
+    active = i <= N_spins
+    Mxy_r = zero(T)
+    Mxy_i = zero(T)
+    Mz = zero(T)
+    ρ = zero(T)
+    ΔBz = zero(T)
+    T1 = T(1)
+    T2 = T(1)
+    x = zero(T)
+    y = zero(T)
+    z = zero(T)
+    Bx_prev = zero(T)
+    By_prev = zero(T)
+    Bz_prev = zero(T)
+
+    if active
         ΔBz = p_ΔBz[i]
         Mxy_r, Mxy_i = reim(M_xy[i])
         Mz = M_z[i]
@@ -38,10 +51,12 @@
         x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, 1)
         Bx_prev, By_prev = reim(s_B1[1])
         Bz_prev = x * s_Gx[1] + y * s_Gy[1] + z * s_Gz[1] + ΔBz - s_Δf[1] / T(γ)
+    end
 
-        ADC_idx = 1u32
-        s_idx = 2u32
-        while (s_idx <= s_length)
+    ADC_idx = 1u32
+    s_idx = 2u32
+    while (s_idx <= s_length)
+        if active
             if MOTION
                 x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, s_idx)
             end
@@ -84,19 +99,22 @@
             Mxy_i = Mxy_new_i * E2
             Mz = Mz_new * E1 + ρ * (T(1) - E1)
 
-            # Acquire Signal
-            if s_idx <= s_length && s_ADC[s_idx]
-                sig_r, sig_i = reduce_signal!(Mxy_r, Mxy_i, sig_group_r, sig_group_i, i_l, N, T, Val(USE_WARP_REDUCTION))
-                if i_l == 1u32
-                    sig_output[i_g, ADC_idx] = complex(sig_r, sig_i)
-                end
-                ADC_idx += 1u32
-            end
-
             Bx_prev, By_prev, Bz_prev = Bx_next, By_next, Bz_next
-            s_idx += 1u32
         end
 
+        # Acquire Signal
+        if HAS_ADC && s_ADC[s_idx]
+            sig_r, sig_i = reduce_signal!(Mxy_r, Mxy_i, sig_group_r, sig_group_i, i_l, N, T, Val(USE_WARP_REDUCTION))
+            if i_l == 1u32
+                sig_output[i_g, ADC_idx] = complex(sig_r, sig_i)
+            end
+            ADC_idx += 1u32
+        end
+
+        s_idx += 1u32
+    end
+
+    if active
         # RF frame -> Rotating frame
         # M * exp(i * ψ)
         ψ_end = s_ψ[s_length]

KomaMRICore/src/simulation/SimMethods/Bloch/gpu/PrecessionKernel.jl

@@ -5,20 +5,19 @@
     M_xy, M_z, 
     @Const(p_x), @Const(p_y), @Const(p_z), @Const(p_ΔBz), @Const(p_T1), @Const(p_T2), @Const(p_ρ), N_spins,
     @Const(s_Gx), @Const(s_Gy), @Const(s_Gz), @Const(s_Δt), @Const(s_ADC), s_length,
-    ::Val{MOTION}, ::Val{USE_WARP_REDUCTION},
+    ::Val{MOTION}, ::Val{USE_WARP_REDUCTION}, ::Val{HAS_ADC},
     sim_method::BlochLikeSimMethods
-) where {T, MOTION, USE_WARP_REDUCTION}
+) where {T, MOTION, USE_WARP_REDUCTION, HAS_ADC}
 
     @uniform N = @groupsize()[1]
     i_l = @index(Local, Linear)
     i_g = @index(Group, Linear)
     i = (i_g - 1u32) * UInt32(N) + i_l
 
-    sig_group_r = @localmem T USE_WARP_REDUCTION ? 32 : N
-    sig_group_i = @localmem T USE_WARP_REDUCTION ? 32 : N
-    sig_r = zero(T)
-    sig_i = zero(T)
+    sig_group_r = @localmem T HAS_ADC ? (USE_WARP_REDUCTION ? 32 : N) : 1
+    sig_group_i = @localmem T HAS_ADC ? (USE_WARP_REDUCTION ? 32 : N) : 1
 
+    active = i <= N_spins
     Mxy_r = zero(T)
     Mxy_i = zero(T)
     t = zero(T)
@@ -32,10 +31,8 @@
     Bz_next = zero(T)
 
     # Setting per-thread (ith spin) properties, and Bz for t0
-    if i <= N_spins
+    if active
         Mxy_r, Mxy_i = reim(M_xy[i])
-        sig_r = Mxy_r
-        sig_i = Mxy_i
         ΔBz = p_ΔBz[i]
         T2 = p_T2[i]
         x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, 1)
@@ -45,7 +42,7 @@
     ADC_idx = 1u32
     s_idx = 2u32
     while s_idx <= s_length
-        if i <= N_spins
+        if active
             if MOTION
                 x, y, z = get_spin_coordinates(p_x, p_y, p_z, i, s_idx)
             end
@@ -56,8 +53,10 @@
             ϕ += (Bz_prev + Bz_next) * T(-π * γ) * Δt
         end
         # Acquire Signal
-        if s_idx <= s_length && s_ADC[s_idx]
-            if i <= N_spins
+        if HAS_ADC && s_ADC[s_idx]
+            sig_r = zero(T)
+            sig_i = zero(T)
+            if active
                 E2 = exp(-t / T2)
                 cis_ϕ_i, cis_ϕ_r = sincos(ϕ)
                 sig_r = E2 * (Mxy_r * cis_ϕ_r - Mxy_i * cis_ϕ_i)
@@ -74,7 +73,7 @@
         s_idx += 1u32
     end
     # Save magnetization at end of block
-    if i <= N_spins
+    if active
         E1 = exp(-t / p_T1[i])
         E2 = exp(-t / T2)
         cis_ϕ_i, cis_ϕ_r = sincos(ϕ)
@@ -83,4 +82,4 @@
     end
 end
 
-## COV_EXCL_STOP
+## COV_EXCL_STOP

KomaMRICore/test/Project.toml

@@ -3,7 +3,6 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KomaMRIBase = "d0bc0b20-b151-4d03-b2a4-6ca51751cb9c"
 KomaMRICore = "4baa4f4d-2ae9-40db-8331-a7d1080e3f4e"
-Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"