Allocates when extracting Dislocation creep parameters from the database

[boriskaus]

This summarises a Discord discussion with @albert-de-montserrat.

Currently, we have allocations when we extract a nonlinear rheology from the Dislocation creep database as in:

function main()
    # Unit system
    CharDim    = SI_units(length=1000m, temperature=1000C, stress=1e7Pa, viscosity=1e20Pas)
    # Numerical parameters
    Ncy        = 100
    # Allocate arrays
    ε0         =    1e-4
    Tv         =    rand(Ncy+1)
    ε̇ii        = ε0*ones(Ncy+1) 
    η          =   zeros(Ncy+1)

    # Configure viscosity model
    flow_nd = SetDislocationCreep("Diabase | Caristan (1982)",CharDim)

    # Setup up viscosity model
    a = @allocated begin
        for i in eachindex(ε̇ii)
            η[i] = compute_viscosity_εII(flow_nd, ε̇ii[i], (;T=Tv[i]))
        end
    end
    @show a
end

julia> main();
a = 4848

If we, however, define the same rheology directly within the routine it does not allocate:

function main1()
    # Unit system
    CharDim    = SI_units(length=1000m, temperature=1000C, stress=1e7Pa, viscosity=1e20Pas)
    # Numerical parameters
    Ncy        = 100
    # Allocate arrays
    ε0         =    1e-4
    Tv         =    rand(Ncy+1)
    ε̇ii        = ε0*ones(Ncy+1) 
    η          =   zeros(Ncy+1)

    # Configure viscosity model
    flow_nd0 = DislocationCreep(;
        Name = "Diabase | Caristan (1982)",
        n = 3.05NoUnits,
        A = 6.0e-2MPa^(-61 // 20) / s,
        E = 276kJ / mol,
        V = 0m^3 / mol,
        r = 0NoUnits,
        Apparatus = AxialCompression,
    )
    flow_nd  = Transform_DislocationCreep(flow_nd0, CharDim)

    # Setup up viscosity model
    a = @allocated begin
        for i in eachindex(ε̇ii)
            η[i] = compute_viscosity_εII(flow_nd, ε̇ii[i], (;T=Tv[i]))
        end
    end
    @show a
end

julia> main1()
a = 0

The underlying problem appears to be that the compiler does not know the type of the units for A at compile-time, as it has a power law exponent in them (here MPa^(-61 // 20) / s), which is different for different types of creep rheologies.
The same issue will likely exist in other types of rheologies as well (Peierls creep, for example).

Two possible solutions suggested by @albert-de-montserrat are:

1. Do not use a Dict as database, but instead reimplement all rheologies to use multiple dispatch. If we want to keep names of the creep laws that have spaces/special symbols in them, the user would have to write something like SetDislocationCreep(Val(Symbol("CreepLawName"))). The alternative option is to change the names to be one word, so the above could become: SetDislocationCreep(:Diabase_Caristan_1982).
2. Try to use TypeSortedCollections.jl. Disadvantage is that it increases the compilation time.

[albert-de-montserrat]

Regarding option 1. Using Symbol instead of String would reduce some typing, but it still requires to be called using Val, i.e. SetDislocationCreep(Val(:Diabase_Caristan_1982)). More annoyingly, we could also make a Diabase_Caristan_1982 <: DislocationCreep type...

[albert-de-montserrat]

I must clarify something I have realised just now. The non-allocating MWE non-allocates only in Julia 1.10, it does allocate in 1.9. Can you reproduce the allocations?

[boriskaus]

I must clarify something I have realised just now. The non-allocating MWE non-allocates only in Julia 1.10, it does allocate in 1.9. Can you reproduce the allocations?

I reproduced the examples above on my machine using 1.9. the # of allocations is less than what you reported

[albert-de-montserrat]

Another light-weight option (works in 1.10) is to define the structs within its own named functions.

function DryOlivine_HirthKohlstedt_2003()
        DislocationCreep(;
            Name = "Dry Olivine | Hirth & Kohlstedt (2003)",
            n = 3.5NoUnits,
            r = 0.0NoUnits,
            A = 1.1e5MPa^(-7 // 2) / s,
            E = 530.0kJ / mol,
            V = 14e-6m^3 / mol,
            Apparatus = AxialCompression,
        )
end

function WetOlivine_HirthKohlstedt_2003b()
        DislocationCreep(;
            Name = "2. Wet Olivine | Hirth & Kohlstedt (2003)",
            n = 3.0NoUnits,
            A = 1600MPa^(-3) / s,
            E = 520.0kJ / mol,
            V = 22.0m^3 / mol,
            r = 1.2NoUnits,
            Apparatus = AxialCompression,
        )
end

function QuartzDiorite_Hansen_Carter_1982()
        DislocationCreep(;
            Name = "Quartz Diorite | Hansen & Carter (1982)",
            n = 2.25NoUnits,
            A = 3.5e-2MPa^(-9 // 4) / s,
            E = 212kJ / mol,
            V = 0m^3 / mol,
            r = 0NoUnits,
            Apparatus = AxialCompression,
        )
end

dislocation_database(f::F) where F = f()

function main()
    p = dislocation_database(QuartzDiorite_Hansen_Carter_1982)
    @allocated compute_viscosity_εII(p, 1.0, (;))
end

and

julia> main()
0

[albert-de-montserrat]

I must clarify something I have realised just now. The non-allocating MWE non-allocates only in Julia 1.10, it does allocate in 1.9. Can you reproduce the allocations?

I reproduced the examples above on my machine using 1.9. the # of allocations is less than what you reported

Weird.

julia> main1()
a = 4848
4848

julia> versioninfo()
Julia Version 1.9.3
Commit bed2cd540a1 (2023-08-24 14:43 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 10 × Apple M2 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, apple-m1)
  Threads: 4 on 6 virtual cores
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 4

[boriskaus]

That would work too - is there an easy way to list all available options in this case? And what about 1.9 with this method?

[albert-de-montserrat]

That would work too - is there an easy way to list all available options in this case?

I left the original databases in the Dictionary in a Data_deprecated folder. I guess It can still be used to list all the rheologies. Perhaps we can also provide another Dict that maps the String name to the function name. Could be useful to get the info from the REPL or for the docs.

This method still allocates in 1.9. Haven't looked in detail of what has changed.

[albert-de-montserrat]

@code_warntype main1() in 1.10:

  flow_nd::DiffusionCreep{Float64, 37, Unitful.FreeUnits{, NoDims, nothing}, Unitful.FreeUnits{(μm³·⁰, MPa⁻¹·⁰, s⁻¹·⁰), 𝐋⁴·⁰ 𝐓 𝐌⁻¹·⁰, nothing}, Unitful.FreeUnits{(kJ, mol⁻¹·⁰), 𝐋²·⁰ 𝐌 𝐍⁻¹·⁰ 𝐓⁻²·⁰, nothing}, Unitful.FreeUnits{(m³·⁰, mol⁻¹·⁰), 𝐋³·⁰ 𝐍⁻¹·⁰, nothing}, Unitful.FreeUnits{(J, K⁻¹·⁰, mol⁻¹·⁰), 𝐋²·⁰ 𝐌 𝐍⁻¹·⁰ 𝚯⁻¹·⁰ 𝐓⁻²·⁰, nothing}} 

and in 1.9

%62 = flow_nd::DiffusionCreep{Float64, _A, Unitful.FreeUnits{, NoDims, nothing}, Unitful.FreeUnits{(μm³·⁰, MPa⁻¹·⁰, s⁻¹·⁰), 𝐋⁴·⁰ 𝐓 𝐌⁻¹·⁰, nothing}, Unitful.FreeUnits{(kJ, mol⁻¹·⁰), 𝐋²·⁰ 𝐌 𝐍⁻¹·⁰ 𝐓⁻²·⁰, nothing}, Unitful.FreeUnits{(m³·⁰, mol⁻¹·⁰), 𝐋³·⁰ 𝐍⁻¹·⁰, nothing}, Unitful.FreeUnits{(J, K⁻¹·⁰, mol⁻¹·⁰), 𝐋²·⁰ 𝐌 𝐍⁻¹·⁰ 𝚯⁻¹·⁰ 𝐓⁻²·⁰, nothing}} where _A

So in 1.10 it is able to infer the length of the Name as a tuple of 37 Chars while 1.9 is not able to infer it....

[boriskaus]

What if we change this in the structure definition:

N = length(Name)

to

N = Int64(length(Name))

[boriskaus]

I must clarify something I have realised just now. The non-allocating MWE non-allocates only in Julia 1.10, it does allocate in 1.9. Can you reproduce the allocations?

I reproduced the examples above on my machine using 1.9. the # of allocations is less than what you reported

Weird.

julia> main1()
a = 4848
4848

julia> versioninfo()
Julia Version 1.9.3
Commit bed2cd540a1 (2023-08-24 14:43 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 10 × Apple M2 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, apple-m1)
  Threads: 4 on 6 virtual cores
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 4

I can confirm that I get the same now after restarting Julia. Mysterious...

[boriskaus]

Perhaps this is caused by the following lines:

function str2tuple(x::String) 
    N = length(x)
    ntuple(i -> x[i], Val(N))
end

on 1.9:

julia> @code_warntype str2tuple("tes")
MethodInstance for GeoParams.MaterialParameters.ConstitutiveRelationships.str2tuple(::String)
  from str2tuple(x::String) @ Main ~/.julia/dev/GeoParams/src/CreepLaw/CreepLaw.jl:64
Arguments
  #self#::Core.Const(GeoParams.MaterialParameters.ConstitutiveRelationships.str2tuple)
  x::String
Locals
  #15::var"#15#16"{String}
  N::Int64
Body::Tuple{Vararg{Char}}
1 ─      (N = Main.length(x))
│   %2 = Main.:(var"#15#16")::Core.Const(var"#15#16")
│   %3 = Core.typeof(x)::Core.Const(String)
│   %4 = Core.apply_type(%2, %3)::Core.Const(var"#15#16"{String})
│        (#15 = %new(%4, x))
│   %6 = #15::var"#15#16"{String}
│   %7 = Main.ntuple(%6, N)::Tuple{Vararg{Char}}
└──      return %7

[albert-de-montserrat]

Yes :) that's why I wanted to remove Name from the structs. The length of a string is unknown at compile time (to make Val(N) type stable, N has to be a literal or a type known at compile time). I was actually surprised that this worked in 1.10

[boriskaus]

well the point is that for all predefined rheologies in our database, we actually do know the length of the Name string when we compile GeoParams. Can that not be put to good use?

[albert-de-montserrat]

A bit annoying but we could directly pass the names as Name=('a','b','c') and maybe change the default kwarg Name=(''). Then it will only allocate if the user actually passes a string, as we would need to use str2char as a fallback.

[boriskaus]

was trying that using str2tuple but did not get that working in an allocation-free manner

[albert-de-montserrat]

It will always allocate, it's like trying to convert the x heap allocated array n = 20; x=rand(n) into a tuple.

In order not to allocate unrolling is needed, and for that to happen we need to know n at compile time (i.e. has to be a literal).

[albert-de-montserrat]

Ugly, but this exists https://github.com/mkitti/StaticStrings.jl . We could use it for the database at least

[boriskaus]

and if we use StaticArrays instead?

[albert-de-montserrat]

I dont think it works for strings.

[boriskaus]

What StaticStrings essentially does is this, I believe:

function str2tuple(x::String) 
    N = 10
    x = rpad(x,N)
    return ntuple(i -> x[i], Val(N))
end

That is type-stable:

julia> @code_warntype str2tuple("test")
MethodInstance for GeoParams.MaterialParameters.ConstitutiveRelationships.str2tuple(::String)
  from str2tuple(x::String) @ Main ~/.julia/dev/GeoParams/src/CreepLaw/CreepLaw.jl:64
Arguments
  #self#::Core.Const(GeoParams.MaterialParameters.ConstitutiveRelationships.str2tuple)
  x@_2::String
Locals
  #32::var"#32#33"{String}
  N::Int64
  x@_5::String
Body::NTuple{10, Char}
1 ─       (x@_5 = x@_2)
│         (N = 10)
│         (x@_5 = Main.rpad(x@_5, N::Core.Const(10)))
│   %4  = Main.:(var"#32#33")::Core.Const(var"#32#33")
│   %5  = Core.typeof(x@_5)::Core.Const(String)
│   %6  = Core.apply_type(%4, %5)::Core.Const(var"#32#33"{String})
│         (#32 = %new(%6, x@_5))
│   %8  = #32::var"#32#33"{String}
│   %9  = Main.Val(N::Core.Const(10))::Core.Const(Val{10}())
│   %10 = Main.ntuple(%8, %9)::NTuple{10, Char}
└──       return %10

I agree that this is a bit ugly; we would have to set N=100 or so, and spit an error message if the input string is longer

[albert-de-montserrat]

That's stable because it's doing padding to a known length of 10. I guess we could force a long enough length for the names. Then need to make them pretty again, otherwise:

julia> x = str2tuple("tes")
('t', 'e', 's', ' ', ' ', ' ', ' ', ' ', ' ', ' ')

julia> str = collect(x) |> String
"tes       "

[boriskaus]

With this change:

"""
    str2tuple(x::String) 
Converts a string to a tuple with fixed length
"""
function str2tuple(x::String) 
    N = 100
    if length(x)>N
        error("Name String is too long; max. allowed length=$N")
    end
    x = rpad(x,N)
    return ntuple(i -> x[i], Val(N))
end

function main1()
    # Unit system
    CharDim    = SI_units(length=1000m, temperature=1000C, stress=1e7Pa, viscosity=1e20Pas)
    # Numerical parameters
    Ncy        = 100
    # Allocate arrays
    ε0         =    1e-4
    Tv         =    rand(Ncy+1)
    ε̇ii        = ε0*ones(Ncy+1) 
    η          =   zeros(Ncy+1)

    # Configure viscosity model
    flow_nd0 = DislocationCreep(;
        Name = str2tuple("Diabase | Caristan (1982)"),
        n = 3.05NoUnits,
        A = 6.0e-2MPa^(-61 // 20) / s,
        E = 276kJ / mol,
        V = 0m^3 / mol,
        r = 0NoUnits,
        Apparatus = AxialCompression,
    )
    flow_nd  = Transform_DislocationCreep(flow_nd0, CharDim)

    # Setup up viscosity model
    a = @allocated begin
        for i in eachindex(ε̇ii)
            η[i] = compute_viscosity_εII(flow_nd, ε̇ii[i], (;T=Tv[i]))
        end
    end
    @show a
end

julia> main1();
a = 0

[boriskaus]

Then need to make them pretty again, otherwise:

Sure, which can be done with strip as in:

function show(io::IO, g::DislocationCreep)
    return print(
        io,
        "DislocationCreep: Name = $(strip(String(collect(g.Name)))), n=$(Value(g.n)), r=$(Value(g.r)), A=$(Value(g.A)), E=$(Value(g.E)), V=$(Value(g.V)), FT=$(g.FT), FE=$(g.FE), Apparatus=$(g.Apparatus)",
    )
end

We do send more data to/from the GPU in this case; I wonder whether that has a significant impact on the performance or memory consumption.

[boriskaus]

So with this change, main1() works in an allocation-free manner on 1.9, which is now added to this branch.

It does, however, not resolve the original issue which is that main() allocates.

[boriskaus]

One potential solution could be by defining the database entries as a tuple (along with the changes discussed above):

"""
    entry = extract_database_entry(Name::String, database::NTuple{N,AbstractCreepLaw})

Extracts an entry from a creeplaw database
"""
function extract_database_entry(Name::String, database::NTuple{N,AbstractCreepLaw}) where {N}
    names = extract_database_names(database)
    found = false
    name_pad = rpad(Name, length(names[1]))
    entry = database[1]
    for i = 1:N
        if (names[i] .==  name_pad)
            entry = database[i]
            found = true
        end
    end
    if !found; error("Unknown database entry: $Name"); end

   return entry
end

"""
    names = extract_database_names(database::Tuple)
Returns a vector with all `names` in the `database`
"""
function extract_database_names(database::Tuple)
    return [String(collect(f.Name)) for f in database]
end

function main()

    flowlaws = ( 
                DislocationCreep(;
                    Name = "Dry Olivine | Hirth & Kohlstedt (2003)",
                    n = 3.5NoUnits,
                    r = 0.0NoUnits,
                    A = 1.1e5MPa^(-7 // 2) / s,
                    E = 530.0kJ / mol,
                    V = 14e-6m^3 / mol,
                    Apparatus = AxialCompression,
                ),
                DislocationCreep(;
                    Name = "2. Wet Olivine | Hirth & Kohlstedt (2003)",
                    n = 3.0NoUnits,
                    A = 1600MPa^(-3) / s,
                    E = 520.0kJ / mol,
                    V = 22.0m^3 / mol,
                    r = 1.2NoUnits,
                    Apparatus = AxialCompression,
                ),
                DislocationCreep(;
                    Name = "Diabase | Caristan (1982)",
                    n = 3.05NoUnits,
                    A = 6.0e-2MPa^(-61 // 20) / s,
                    E = 276kJ / mol,
                    V = 0m^3 / mol,
                    r = 0NoUnits,
                    Apparatus = AxialCompression,
                )
    )

    # Unit system
    CharDim    = SI_units(length=1000m, temperature=1000C, stress=1e7Pa, viscosity=1e20Pas)
    # Numerical parameters
    Ncy        = 100
    # Allocate arrays
    ε0         =    1e-4
    Tv         =    rand(Ncy+1)
    ε̇ii        = ε0*ones(Ncy+1) 
    η          =   zeros(Ncy+1)

    # Configure viscosity model
    flow_nd0 = extract_database_entry("Diabase | Caristan (1982)", flowlaws)
    flow_nd  = Transform_DislocationCreep(flow_nd0, CharDim)

    # Setup up viscosity model
    a = @allocated begin
        for i in eachindex(ε̇ii)
            η[i] = compute_viscosity_εII(flow_nd, ε̇ii[i], (;T=Tv[i]))
        end
    end
    return a
end

julia> main()
0

[albert-de-montserrat]

That was one of my attempts. Works for small tuples, but breaks down if you add the whole database:

julia> main() # first call
42368

julia> main() # second call
4848

I don't dislike the idea of wrapping them into their own functions -which makes them already a type.

[boriskaus]

This can now be closed with the merge of PR #144