Rework the static option to make a fully isbits flash that can run on GPU (breaking?)

Project.toml

@@ -12,6 +12,8 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 ForwardDiff = "0.10, 1"
+JLArrays = "0.1, 0.2, 0.3"
+KernelAbstractions = "0.9"
 LinearAlgebra = "1"
 Printf = "1"
 Roots = "1, 2, 3"
@@ -20,8 +22,10 @@ julia = "1.6"
 
 [extras]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "ForwardDiff", "StaticArrays"]
+test = ["Test", "ForwardDiff", "JLArrays", "KernelAbstractions", "StaticArrays"]

docs/src/api/flash.md

@@ -13,7 +13,7 @@ Private = false
 ## Two-phase flash
 ```@autodocs
 Modules = [MultiComponentFlash]
-Pages   = ["flash.jl", "flash_types.jl"]
+Pages   = ["flash.jl", "flash_types.jl", "static.jl"]
 Order   = [:type, :function]
 Private = false
 ```

docs/src/examples/advanced.md

@@ -38,64 +38,49 @@ S = flash_storage(eos, conditions, method = m)
 16
 ```
 
-See the unit tests for examples where the flash can use `StaticArrays` to avoid allocations entirely.
+## Immutable performance and GPU use
 
-## Performance example
+[`flash_2ph_immutable`](@ref) is the public interface to the fully static SSI path.
+It is useful when the component count is small and fixed, particularly inside CPU
+or GPU kernels. Convert the EOS once with [`make_eos_immutable`](@ref), and provide the
+overall composition as an `SVector`:
 
-The default interface is designed for ease-of-use with standard Julia types, but the module  also supports further by using `StaticArrays`:
+```julia
+using BenchmarkTools, MultiComponentFlash, StaticArrays
+
+eos_static = make_eos_immutable(eos)
+conditions_static = (p = p, T = T, z = SVector{length(z)}(z))
+
+V, K = flash_2ph_immutable(eos_static, conditions_static)
+@btime flash_2ph_immutable($eos_static, $conditions_static)
+```
+
+`V` is the scalar vapor fraction and `K` is an `SVector`. All working vectors are
+immutable values local to the call; the input `conditions_static.z` must also be an
+`SVector`. The implementation currently supports `GenericCubicEOS` with
+`SSIFlash`, and compilation is specialized on the number of components.
+
+The two-argument form creates the static storage marker automatically. It can also
+be constructed once and passed as the final positional argument:
 
 ```julia
-using MultiComponentFlash, BenchmarkTools, StaticArrays
-function bench(m, static_size = false)
-    p = 6e6
-    T = 480.0
-    # Take the SPE5 benchmark
-    eos, data = cubic_benchmark("spe5")
-    n = number_of_components(eos)
-    z = repeat([1/n], n)
-    conditions = (p = p, T = T, z = z)
-    S = flash_storage(eos, conditions, method = m, static_size = static_size)
-    K = initial_guess_K(eos, conditions)
-    if static_size
-        N = number_of_components(eos)
-        K = MVector{N}(K)
-    end
-    V, K, status = flash_2ph!(S, K, eos, conditions, NaN, method = m, extra_out = true)
-    println("V = $V (Completed in $(status.its) iterations)")
-    @btime flash_2ph!($S, $K, $eos, $conditions, NaN, method = $m)
-    return nothing
-end
-println("SSI:")
-bench(SSIFlash())
-println("SSI (static arrays):")
-bench(SSIFlash(), true)
-##
-println("Newton:")
-bench(NewtonFlash())
-println("Newton (static arrays):")
-bench(NewtonFlash(), true)
+storage = flash_storage(eos_static, conditions_static; static = true)
+V, K = flash_2ph_immutable(eos_static, conditions_static, storage)
+@btime flash_2ph_immutable($eos_static, $conditions_static, $storage)
 ```
 
-The output will be a bit different on other CPUs, but this flash generally takes around 20 microseconds to complete, including both stability test and flash.
+Static storage is a zero-size immutable marker rather than a mutable work buffer,
+so constructing it inline normally compiles away and does not allocate. Passing it
+explicitly can still be convenient when setting up a kernel. For example, each
+kernel work item can construct its conditions and call:
 
 ```julia
-SSI:
-V = 0.03279769425318795 (Completed in 14 iterations)
-  18.500 μs (0 allocations: 0 bytes)
-SSI (static arrays):
-V = 0.03279769425318795 (Completed in 14 iterations)
-  16.500 μs (0 allocations: 0 bytes)
-
-Newton:
-V = 0.032797694260046494 (Completed in 4 iterations)
-  20.100 μs (0 allocations: 0 bytes)
-Newton (static arrays):
-V = 0.032797694260046494 (Completed in 4 iterations)
-  19.900 μs (0 allocations: 0 bytes)
+conditions_i = (p = pressure[i], T = temperature[i], z = z_static)
+V, K = flash_2ph_immutable(eos_static, conditions_i, storage)
 ```
 
-!!! note "Use of `StaticArrays`"
-    Switching to statically sized arrays can improve the speed, at the cost of longer compilation times. Please note that for `StaticArrays` there will be compilation that is dependent on the number of components in your mixture. For example, switching from a five to six component mixture will trigger a full recompilation of your chosen flash.
+Do not share ordinary mutable storage from `flash_storage(...; static = false)`
+between kernel work items.
 
 ## Generate and plot a phase diagram
 

src/MultiComponentFlash.jl

@@ -19,7 +19,7 @@ module MultiComponentFlash
     export KValuesEOS
     export number_of_components
     # Flash interfaces
-    export flash_2ph, flash_2ph!, flash_storage
+    export flash_2ph, flash_2ph!, flash_2ph_immutable, flash_storage
     export stability_2ph, stability_2ph!
     # Algorithms for flash
     export SSIFlash, NewtonFlash, SSINewtonFlash
@@ -35,6 +35,7 @@ module MultiComponentFlash
 
     export force_scalars, force_coefficients, force_coefficients!
     export critical_pressure, critical_temperature, critical_volume, acentric_factor, molar_weight
+    export make_eos_immutable
     export cubic_benchmark
     export single_phase_label
 
@@ -53,6 +54,7 @@ module MultiComponentFlash
     include("flash.jl")
     include("derivatives.jl")
     include("stability.jl")
+    include("static.jl")
     include("tables.jl")
 
     include("flow_coupler.jl")

src/eos.jl

@@ -11,9 +11,11 @@ number_of_components(e::AbstractEOS) = number_of_components(e.mixture)
 forces_per_phase(eos::GenericCubicEOS) = false
 
 function get_phase(cond)
-    return get(cond, :phase, :unknown)::Symbol
+    return phase_symbol(get(cond, :phase, :unknown))
 end
 
+@inline phase_symbol(phase::Symbol) = phase
+
 function set_phase(cond, phase::Symbol, throw::Bool = false)
     if throw && haskey(cond, :phase) && cond.phase != :unknown
         throw(ArgumentError("Phase state already set to $(cond.phase), cannot change to $phase."))
@@ -93,34 +95,40 @@ minimum_allowable_root(eos, forces, scalars) = 1e-16
     return roots
 end
 
-function pick_root(eos, roots, cond, forces, scalars)
-    phase = get_phase(cond)
-    r_ϵ = minimum_allowable_root(eos, forces, scalars)
-    max_r = maximum(roots)
-    min_r = minimum((x) -> x > r_ϵ ? x : Inf, roots)
-    if min_r == max_r
-        r = min_r
-    elseif phase == :liquid
-        r = min_r
-    elseif phase == :vapor
-        r = max_r
-    else
-        function Gibbs(Z)
-            E = 0.0
-            z = cond.z
-            @inbounds for i in eachindex(z)
-                ϕ = component_fugacity_coefficient(eos, cond, i, Z, forces, scalars)
-                E += z[i]*ϕ
-            end
-            return E
+@inline function root_bounds(roots, minimum_root)
+    max_root = -Inf
+    min_root = Inf
+    for root in roots
+        max_root = max(max_root, root)
+        if root > minimum_root
+            min_root = min(min_root, root)
         end
-        if Gibbs(min_r) < Gibbs(max_r)
-            r = min_r
-        else
-            r = max_r
+    end
+    return min_root, max_root
+end
+
+@inline function pick_root(eos, roots, cond, forces, scalars)
+    phase = get(cond, :phase, :unknown)
+    return pick_root(eos, roots, cond, forces, scalars, phase)
+end
+
+function pick_root(eos, roots, cond, forces, scalars, phase::Symbol)
+    min_r, max_r = root_bounds(roots, minimum_allowable_root(eos, forces, scalars))
+    if min_r == max_r || phase == :liquid
+        return min_r
+    elseif phase == :vapor
+        return max_r
+    end
+    function Gibbs(Z)
+        E = 0.0
+        z = cond.z
+        @inbounds for i in eachindex(z)
+            ϕ = component_fugacity_coefficient(eos, cond, i, Z, forces, scalars)
+            E += z[i]*ϕ
         end
+        return E
     end
-    return r
+    return Gibbs(min_r) < Gibbs(max_r) ? min_r : max_r
 end
 
 """
@@ -138,9 +146,9 @@ function force_coefficients(eos::AbstractCubicEOS, cond; static_size = false)
     n = number_of_components(eos)
     eT = Base.promote_eltype(cond.p, cond.T, cond.z[1])
     if static_size
-        A_ij = @MMatrix zeros(eT, n, n)
-        A_i = @MVector zeros(eT, n)
-        B_i = @MVector zeros(eT, n)
+        A_ij = zero(MMatrix{n, n, eT})
+        A_i = zero(MVector{n, eT})
+        B_i = zero(MVector{n, eT})
     else
         A_ij = zeros(eT, n, n)
         A_i = zeros(eT, n)

src/eos_types.jl

@@ -11,9 +11,9 @@ definitions for the terms (they are, after all, all cubic in form). References:
  2. [Simulation of Gas Condensate Reservoir Performance  by K.H. Coats](https://doi.org/10.2118/10512-PA)
 
 """
-struct GenericCubicEOS{T, R, N, V} <: AbstractCubicEOS
+struct GenericCubicEOS{T, R, N, V, M<:MultiComponentMixture{R, N}} <: AbstractCubicEOS
     type::T
-    mixture::MultiComponentMixture{R, N}
+    mixture::M
     m_1::R
     m_2::R
     ω_a::R