Use value types to store model parameters

[Mikolaj-A-Kowalski]

This issue is to track/investigate potential improvement opportunity by allowing compiler to specialise on the values on model parameters.

At the moment, various models (e.g. SugarKelp) use a large number of parameters, which are stored as fields in an immutable struct. The performance problem of this approach is that they are not technically compile-time constants and compiler must generate code that will support multiple values. As the result, the values need to be loaded (and consume) GPU registers and expressions like exp(param_1)/param_2 runtime evaluated. If the values are instead made type parameters, compiler will be free to pre-evaluate expressions involving them and keep the constants in instruction, freeing GPU registers.

Generally julia discourages doing something like that, but in this case I think it makes sense since we expect parameters to be set only once at the start of a long simulation, hence cases where we would pay a price of extra compile time when changing the parameter values should be rare.

First thing will be to try brutally hardcoding parameters for some model to see what kind of improvement we can get. Then, if it is indeed favorable, we will need to make/find a macro that will make making structs that are a pack of Value types easy.

[Mikolaj-A-Kowalski]

I have made an exercise of hardcoding all the parameters of the SugarKelp model as it is on the current main.

The conclusion are mixed :-/ As expected, the register usage has decreased significantly (for example for the tendency calculation of the N the usage dropped from 255 to 205, but since it is not enough to improve occupancy, the runtime did not change (decreased by ~2.18%))

I attach the NCU reports reports.tar.gz and the script LOBSTER_3D_profile_particles.jl.txt.

Even with this results I would argue this direction is worth pursuing as, in principle, it should reduce 'starting point' for any effort to reduce register usage and increase occupancy. 50 less registers is significant, but in this case sadly not enough to push the kernel over the threshold

[jagoosw]

Could we see what the model code looks like when you do this?

[Mikolaj-A-Kowalski]

Could we see what the model code looks like when you do this?

You mean the machine code or the source? For the latter at the moment I did the test just by moving the constants to submodule and making them const. It is obviously not the production solution.

To include it in the code would require having types like this:

get_value_from_value(::Val{V}) where {V} = V

# Macro must produce the following
struct Params{V1, V2, V3}
    val1::Val{V1}
    val2::Val{V2}
    val3::Val{V3}
    Params(v1,v2,v3) = new{v1, v2, v3}()
end

# It is important to use build-in get field on the RHS
# otherwise we end up with infinite recursion
Base.getproperty(p::Params, s::Symbol) = get_value_from_value(Base.getfield(p,s))

but it is still to verbose so I am now looking into producing something like that using a macro.

[glwagner]

Can you Val the whole struct? Eg Val(params)?

[Mikolaj-A-Kowalski]

Can you Val the whole struct? Eg Val(params)?

I have tried with NamedTuple instead of a struct and it worked (i.e. compiled and was correct), but the code as examined by @code_llvm or @code_native did not seem to make it zero-cost abstraction.

Basically:

get_value_from_value(::Val{V}) where {V} = V

struct Params1{V1}
    val1::Val{V1}
    Params1(v1) = new{v1}()
end

struct Data
    a::Float64
    b::Float64
    c::Float64
end

a = Params1(Data(1.1, 2.2, 3.3))

@code_llvm get_value_from_value(a.val1).a

if I interpret it correctly, julia seems to construct an instance of the struct somewhere in memory and load the value from memory location.

LLVM output

; Function Signature: getproperty(Main.Data, Symbol)
;  @ Base.jl:49 within `getproperty`
define double @julia_getproperty_1589(ptr nocapture noundef nonnull readonly align 8 dereferenceable(24) %"x::Data", ptr noundef nonnull %"f::Symbol") #0 {
top:
  %0 = call i32 @ijl_field_index(ptr nonnull @"+Main.Data#1593.jit", ptr nonnull %"f::Symbol", i32 0)
  %.not = icmp eq i32 %0, -1
  br i1 %.not, label %err, label %ok

err:                                              ; preds = %top
  call void @ijl_has_no_field_error(ptr nonnull @"jl_sym#Data#1594.jit", ptr nonnull %"f::Symbol")
  unreachable

ok:                                               ; preds = %top
  %1 = zext i32 %0 to i64
  %2 = getelementptr inbounds double, ptr %"x::Data", i64 %1
  %.unbox = load double, ptr %2, align 8
  ret double %.unbox
}

On the other hand the behaviour is as we want if we put values separately into the struct

get_value_from_value(::Val{V}) where {V} = V


# Macro must produce the following
struct Params{V1, V2, V3}
    val1::Val{V1}
    val2::Val{V2}
    val3::Val{V3}
    Params(v1,v2,v3) = new{v1, v2, v3}()
end
a = Params(1.1, 2.2, 3.3)

@code_llvm get_value_from_value(a.val1)

produces

define double @julia_get_value_from_value_1605() #0 {
top:
  ret double 1.100000e+00
}

[Mikolaj-A-Kowalski]

And writing all of this made me realise that it was not a real problem but just an artifact of how I did the introspection...

get_value_from_value(::Val{V}) where {V} = V

struct Params1{V1}
    val1::Val{V1}
    Params1(v1) = new{v1}()
end

struct Data
    a::Float64
    b::Float64
    c::Float64
end

a = Params1(Data(1.1, 2.2, 3.3))

fetch_value(a) = get_value_from_value(a.val1).a

@code_llvm fetch_value(a)

produces the desired

define double @julia_fetch_value_1796() #0 {
top:
  ret double 1.100000e+00
}

Hence indeed it is not necessary to go down into making a macro and just storing a struct with parameters values is enough 🤔

[glwagner]

I think an interface like

extract(::Val{x}} where x = x
extract(x) = x

does seem ideal. This will avoid having to create two copies of every function.

In principle you could write a macro that would act something like

@valdef function f(x, ::Val{y}) where y
    # something
end

which my autoproduce the auxiliary form

function f(x, y)

then you support "optional Val inputs" for parameters at the function definition level. This might reduce the source code you have to change. Either way you do have to change source because you cannot design a macro @valdef f(x, y) without marking which of the arguments x, y need to be val'd.